# AUTOGENERATED! DO NOT EDIT! File to edit: 00_pde.ipynb (unless otherwise specified).

__all__ = ['PDEDataset', 'PDE']

# Cell

from .utils import *
from .solution import *
import torch
from torch.utils.data import Dataset, DataLoader
from fastprogress import master_bar, progress_bar
import numpy as np

class PDEDataset(Dataset):
    "Receives a dict of arrays and returns every possible combination of the elements in the arrays"
    def __init__(self, inputs):
        # convert dict to array
        self.inputs = np.array([inputs[k] for k in inputs])
        # length of the dataset (all possible combinations)
        self.len = 1
        for input in self.inputs:
            self.len *= len(input)
        # modules
        self.mods = []
        for i, _ in enumerate(self.inputs):
            mod = 1
            for j, input in enumerate(self.inputs):
                if j < i:
                    mod *= len(input)
            self.mods.append(mod)

    def __len__(self):
        return self.len

    def __getitem__(self, idx):
        return torch.FloatTensor([inp[(idx // self.mods[i]) % len(inp)] for i, inp in enumerate(self.inputs)])

class PDE:
    "PDE class with basic functionality to solve PDEs with NNs"
    def __init__(self, inputs, outputs, params=None, order=1):

        # check lists of unique strings, non-repeated
        checkIsListOfStr(inputs)
        checkUnique(inputs)
        checkIsListOfStr(outputs)
        checkUnique(outputs)
        checkNoRepeated(inputs, outputs)
        if params:
            checkIsListOfStr(params)
            checkUnique(params)
            checkNoRepeated(inputs, params)
            checkNoRepeated(params, outputs)

        # save keys
        self.input_keys = inputs
        self.output_keys = outputs
        self.param_keys = params

        # initialize values
        self.train_inputs = {k: [] for k in self.input_keys}
        self.test_inputs = {k: [] for k in self.input_keys}
        self.outputs = {k: [] for k in self.output_keys}
        self.params = {}
        if self.param_keys:
            self.params = {k: [] for k in self.param_keys}

        # bocos
        self.bocos = []

        self.init = False
        self.order = order
        self.eval = False

    def summary(self):
        "Print a summary of the PDE inputs, outputs, params and bocos."
        print('inputs (train): ', self.train_inputs)
        print('inputs (test): ', self.test_inputs)
        print('outputs: ', self.outputs)
        if self.params:
            print('params: ', self.params)
        print('bocos: ', [boco.type for boco in self.bocos])
        print('')

    def setValues(self, values, train=True):
        self.eval = not train
        "Set values for inputs and params"
        checkValidDict(values)
        for key in values:
            value = values[key]
            if key in self.input_keys:
                if train:
                    self.train_inputs[key] = value
                else:
                    self.test_inputs[key] = value
            elif key in self.param_keys:
                if train:
                    self.params[key] = value
                else:
                    raise Exception('You cannot set params in test data !')
            elif key in self.output_keys:
                raise Exception('You cannot set values to outputs !')
            else:
                raise Exception('Key '+ key +' not found !')

    def addBoco(self, boco):
        "Add a boco to the list of bocos"
        boco.addBoco(self.input_keys, self.output_keys)
        self.bocos += [boco]

    def bocoSummary(self):
        "Print summary of each boco"
        for boco in self.bocos:
            boco.summary(self.input_keys, self.output_keys, self.param_keys)

    def buildSolution(self, topo):
        "Build an MLP to be the solution to the PDE"
        n_inputs, n_outputs = len(self.input_keys), len(self.output_keys)
        self.solution = Solution(n_inputs, n_outputs, topo['layers'], topo['neurons'], topo['activations'])

    def compile(self, lr=0.001, epochs=100, batch_size=32):
        "Set the required parameters for training"
        self.optimizer = torch.optim.Adam(self.solution.parameters(), lr=lr)
        self.epochs = epochs
        self.bs = batch_size

    def solve(self, device, path):
        "Find a solution to the PDE"
        # initialize dataladers
        if not self.init:
            self.initialize()
        # convert params to tensors
        params = {k: torch.FloatTensor(self.params[k]).to(device) for k in self.params}
        # train loop
        self.solution.to(device)
        best_loss = 1e8
        hist = {'train_loss': [], 'bocos': {boco.name: [] for boco in self.bocos}}
        if self.eval:
            hist['val_loss'] = []
        mb = master_bar(range(1, self.epochs+1))
        for epoch in mb:
            #train
            self.solution.train()
            pde_loss = []
            bocos_loss = {boco.name: [] for boco in self.bocos}
            for inputs in progress_bar(self.dataloader['train'], parent=mb):
                # compute pde solution
                inputs = inputs.to(device)
                inputs.requires_grad = True
                #print(inputs.require_grad)
                outputs = self.solution(inputs)
                # compute gradients of outputs w.r.t. inputs
                grads, _inputs = self.computeGrads(inputs, outputs)
                # compute loss
                loss = self.computePDELoss(grads, _inputs, params).pow(2).mean()
                pde_loss.append(loss.item())
                # compute bocos loss
                # compute loss for all boco points at once (1 batch each b.c, optimize all together)
                for boco in self.bocos:
                    boco_loss = boco.computeLoss(self.solution, device)
                    bocos_loss[boco.name].append(boco_loss.item())
                    loss += boco_loss
                # update weights
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()
                mb.child.comment = f'train loss {np.mean(pde_loss):.5f}'

            pde_total_loss = np.mean(pde_loss)
            bocos_total_loss = 0
            for boco in self.bocos:
                bocos_loss[boco.name] = np.mean(bocos_loss[boco.name])
                hist['bocos'][boco.name].append(bocos_loss[boco.name])
                bocos_total_loss += bocos_loss[boco.name]
            total_loss = pde_total_loss + bocos_total_loss

            hist['train_loss'].append(total_loss)
            info = f'Epoch {epoch}/{self.epochs} Losses {total_loss:.5f} \n PDE  {pde_total_loss:.5f}'
            for boco in self.bocos:
                info += f'\n {boco.name} {bocos_loss[boco.name]:.5f}'

            if self.eval:
                #evaluate
                self.solution.eval()
                val_loss = []
                for inputs in progress_bar(self.dataloader['val'], parent=mb):
                    # compute pde solution
                    inputs = inputs.to(device)
                    inputs.requires_grad = True
                    outputs = self.solution(inputs)
                    # compute gradients of outputs w.r.t. inputs
                    grads, _inputs = self.computeGrads(inputs, outputs)
                    # compute loss
                    loss = self.computePDELoss(grads, _inputs, params).pow(2).mean()
                    val_loss.append(loss.item())
                    mb.child.comment = f'val loss {np.mean(val_loss):.5f}'
                # save model if best loss (this does not include bocos losses !!!)
                val_total_loss = np.mean(val_loss)
                if val_total_loss < best_loss:
                    best_loss = total_loss
                    torch.save(self.solution.state_dict(), path)

                hist['val_loss'].append(val_total_loss)
                info += f'\n Val {val_total_loss:.5f} '

            mb.write(info)
            #mb.first_bar.comment = f'best acc {best_acc:.5f} at epoch {best_e}'

        return hist

    def initialize(self):
        self.dataset = {
            'train': PDEDataset(self.train_inputs),
            'val': PDEDataset(self.test_inputs)
        }
        self.dataloader = {
            'train': DataLoader(self.dataset['train'], batch_size=self.bs, shuffle=True, num_workers=4),
            'val': DataLoader(self.dataset['val'], batch_size=self.bs, shuffle=False, num_workers=4)
        }
        for boco in self.bocos:
            boco.initialize()
        self.init = True

    def computeGrads(self, inputs, outputs):
        # init grads
        grads = {o: {i: [] for i in self.input_keys} for o in self.output_keys}

        # compute first order derivatives
        for i, output in enumerate(self.output_keys):
            _grads = self.computeGrad(outputs[:,i], inputs)
            #print(_grads)
            # save in dict
            for j, inp in enumerate(self.input_keys):
                grads[output][inp] = _grads[:,j]

            # compute higher order derivatives (only works for second order)
            if self.order > 1:
                #order = 1
                #while order < self.order:
                #order += 1

                order = 2
                grads[f'{order}{output}'] = {}
                # compute gradients of gradients
                for j, inp in enumerate(self.input_keys):
                    __grads = self.computeGrad(_grads[:,j], inputs)
                    # save in dict
                    for k, inp2 in enumerate(self.input_keys):
                        grads[f'{order}{output}'][f'{inp}{inp2}'] = __grads[:,k]

        # assign keys to inputs
        _inputs = {inp: inputs[:,i] for i, inp in enumerate(self.input_keys)}

        return grads, _inputs

    def computeGrad(self, outputs, inputs):
        # torch.autograd.grad -> Computes and returns the sum of gradients of outputs w.r.t. the inputs.
        _grads, = torch.autograd.grad(outputs, inputs,
                        grad_outputs=outputs.data.new(outputs.shape).fill_(1),
                        create_graph=True, only_inputs=True)
        return _grads

    def computePDELoss(self, inputs, outputs, params=None):
        print('This function has to be overloaded by a child class!')

    def load_state_dict(self, path):
        self.solution.load_state_dict(torch.load(path))

    def evaluate(self, inputs, device):
        "Evaluate solution"
        checkValidDict(inputs)
        checkDictArray(inputs, self.input_keys)
        # set values of inpenedent vars
        for key in self.input_keys:
            if key in inputs:
                self.test_inputs[key] = inputs[key]
            else:
                raise Exception(key + ' is not an input')
        # build dataset
        dataset = PDEDataset(self.test_inputs)
        outputs = []
        self.solution.eval()
        for i in range(len(dataset)):
            input = dataset[i].to(device)
            outputs.append(self.solution(input).cpu().detach().numpy())
        return np.array(outputs)