#AUTOGENERATED! DO NOT EDIT! File to edit: dev/generate_wave_eqn_data.ipynb (unless otherwise specified).

__all__ = ['get_initial_conditions', 'plot_image', 'Laplacian', 'batchLaplacian', 'run_wave_forward',
           'generate_wave_data', 'WaveDataset', 'wave_loss']

#Cell
import os, sys
from time import time, sleep
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F
import numpy as np
import argparse
import pprint as pp
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.animation as manimati
from matplotlib import animation, rc
from IPython.display import HTML
from .utils import pkl_save, pkl_load

#Cell
def get_initial_conditions(num_sources=1, max_u0=1., grid_size=128):
    u = torch.zeros((grid_size, grid_size), requires_grad=True)
    source_locs = []
    for src in range(num_sources):
        a, b = np.random.randint(0, grid_size, 2)
        #u[grid_size//2, grid_size//2] = 1.
        u[a, b] = max_u0
        source_locs.append((a,b))
    prev_u = torch.zeros((grid_size, grid_size), requires_grad=True)
    u = Variable(u, requires_grad=True)
    prev_u = Variable(prev_u, requires_grad=True)
    return u.clone(), prev_u.clone(), source_locs

#Cell
def plot_image(f, title='wave amplitude'):
    # plt.ion()
    u_mx = f.max()
    fig = plt.figure()
    ax = fig.add_subplot(111)
    plt.title(title)
    cmap = plt.cm.ocean
    img = ax.imshow(f.data, cmap=cmap, vmin=-u_mx, vmax=u_mx)

    fig.colorbar(img, orientation='vertical')
    plt.show()
    return img, fig

#Cell
def Laplacian(u):
    u = u.unsqueeze(0).unsqueeze(0)
    L = np.array([[.5, 1, .5], [1, -6., 1], [.5, 1, .5]], 'float32')
    pad = (L.shape[0]-1)//2
    conv = nn.Conv2d(1, 1, L.shape[0], 1, bias=False, padding=(pad, pad))
    conv.weight.data = torch.tensor(L).unsqueeze(0).unsqueeze(0)
    return conv(u).squeeze(0).squeeze(0)

#Cell
def batchLaplacian(u,device='cpu'):
    u = u.unsqueeze(1)
    L = np.array([[.5, 1, .5], [1, -6., 1], [.5, 1, .5]], 'float32')
    pad = (L.shape[0]-1)//2
    conv = nn.Conv2d(1, 1, L.shape[0], 1, bias=False, padding=(pad, pad)).to(device)
    conv.weight.data = torch.tensor(L).unsqueeze(0).unsqueeze(0).to(device)
    return conv(u).squeeze()

#Cell
def run_wave_forward(u, prev_u, dt, dx, nt, doPlot=False):
    soln = [u]
    if doPlot:
        img, fig = plot_image(u)

    DT_DX_SQ = (dt/dx)**2
    for _ in range(nt-1):
        next_u = DT_DX_SQ*Laplacian(u) + 2*u - prev_u
        prev_u = u
        u = next_u
        soln.append(u)
        # print(u.data.min())
        if doPlot:
            img.set_data(u.data)
            fig.canvas.draw()
            fig.canvas.flush_events()
            plt.pause(.0001)

    return soln

#Cell
def generate_wave_data(num_sources=10,max_u0=1,grid_size=128,nt=200):
    u, u_prev, source_locs = get_initial_conditions(num_sources,max_u0,grid_size)
    # CFL: dt <= C dx
    dt = .1
    dx = .2
    soln = run_wave_forward(u, u_prev, dt, dx, nt)
    return soln, source_locs

#Cell
from torch.utils.data import Dataset, DataLoader
class WaveDataset(Dataset):
    def __init__(self, n_sims=1,transform=None,numToKeep=np.infty):
        if type(n_sims) == int:
            self.sim_params = []
            data = []
            dt = .1
            dx = .2
            t = dt*np.arange(200)
            t = t/t[-1]
            t = t.astype('float32')
            normalizeX = 10.0
            stacked_frames = []
            for k in range(n_sims):
                d = dict()
                d['num_sources'] = 1 #np.random.randint(5,16)
                #self.sim_params.append(d)
                X, source_locs = generate_wave_data(**d)
                source_locs = (np.array(source_locs)/128) - .5
                source_locs = source_locs.astype('float32')
                X = [xi.detach().numpy()/normalizeX for xi in X]
                for idx in range(1,199):
                    sample = (np.stack((X[idx-1],X[idx],X[idx+1])),t[idx],source_locs)
                    stacked_frames.append(sample)
        else:
            stacked_frames = pkl_load(n_sims)

        if len(stacked_frames) > numToKeep:
            A = []
            A = [stacked_frames[i] for i in range(numToKeep)]
            stacked_frames = A

        self.data = stacked_frames
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        X = self.data[idx]
        if self.transform:
            X = self.transform(X)
        return X

#Cell
def wave_loss(u0,u1,u2,dt=.1,dx=.2,device='cpu'):
    '''
    Measures how well u1 satisifies the wave equation. no source term.
    '''
    DT_DX_SQ = (dt/dx)**2
    d = u2 - DT_DX_SQ*batchLaplacian(u1,device) + 2*u1 - u0
    return torch.mean(torch.abs(d))