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

__all__ = ['Boco', 'PeriodicBocoDataset', 'PeriodicBoco', 'DirichletBocoDataset', 'DirichletBoco', 'NeumannBocoDataset',
           'NeumannBoco']

# Cell

class Boco():
    "Base class to work with bocos"
    def __init__(self, name=None):
        self.type = None
        self.name = name
        self.bs = None
        self.dataset = None
        self.DataLoader = None

    def addBoco(self, input_keys, output_keys):
        "This function is called when a PDE adds this boco"
        print('Override this function to add the required values')

    def check(self, inputs, outputs, params):
        print('Override this function to check everything is ok')

    def summary(self):
        print(f'{self.name} summary:')
        print(f'Type: {self.type}')

    def setSolverParams(self, bs):
        self.bs = bs

# Cell

from .utils import *
import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np

class PeriodicBocoDataset(Dataset):
    "Receives two dicts of arrays and returns every possible combination of the elements in each dict by pairs"
    def __init__(self, inputs1, inputs2):
        self.inputs1 = np.array([inputs1[k] for k in inputs1])
        self.inputs2 = np.array([inputs2[k] for k in inputs2])
        # length of the dataset (all possible combinations)
        self.len = 1
        for input in self.inputs1:
            self.len *= len(input)
        # modules
        self.mods = []
        for i, _ in enumerate(self.inputs1):
            mod = 1
            for j, input in enumerate(self.inputs1):
                if j < i:
                    mod *= len(input)
            self.mods.append(mod)

    def __len__(self):
        return self.len

    def __getitem__(self, idx):
        item1, item2 = np.zeros(len(self.inputs1)), np.zeros(len(self.inputs2))
        for i, input in enumerate(self.inputs1):
            item1[i] = input[(idx // self.mods[i]) % len(input)]
        for i, input in enumerate(self.inputs2):
            item2[i] = input[(idx // self.mods[i]) % len(input)]
        return torch.from_numpy(item1).float(), torch.from_numpy(item2).float()

class PeriodicBoco(Boco):
    "Periodic boundary condition: The outputs of both inputs will be enforced to be equal."
    def __init__(self, name, inputs1, inputs2):
        super().__init__(name)
        self.type = 'periodic'

        # check for dict with numpy arrays, same inputs and same length
        checkValidDict(inputs1)
        checkValidDict(inputs2)
        if len(inputs1) != len(inputs2):
            raise Exception('Inputs must have the same length !')
        for k in inputs1:
            if k not in inputs2:
                raise Exception(k + ' must be present in both inputs !')
            if len(inputs1[k]) != len(inputs2[k]):
                raise Exception(k + ' must have same length in both inputs !')

        self.inputs1 = inputs1
        self.inputs2 = inputs2

    def addBoco(self, input_keys, output_keys):
        "This function is called when a PDE adds this boco"
        # check that all inputs are present
        checkDictArray(self.inputs1, input_keys)
        checkDictArray(self.inputs2, input_keys)
        # reorder to match order in pde
        self.inputs1 = {k: self.inputs1[k] for k in input_keys}
        self.inputs2 = {k: self.inputs2[k] for k in input_keys}

    def summary(self, input_keys, output_keys, param_keys):
        super().summary()
        print('Input 1: ', self.inputs1)
        print('Input 2: ', self.inputs2)
        print('')

    def initialize(self):
        self.dataset = PeriodicBocoDataset(self.inputs1, self.inputs2)
        self.loss = torch.nn.MSELoss()
        # batch training in bocos ??
        self.dataloader = DataLoader(self.dataset, batch_size=len(self.dataset), num_workers=4)
        self._inputs1, self._inputs2 = next(iter(self.dataloader))

    def computeLoss(self, model, device):
        inputs1 = self._inputs1.to(device)
        inputs2 = self._inputs2.to(device)
        outputs1 = model(inputs1)
        outputs2 = model(inputs2)
        return self.loss(outputs1, outputs2)

# Cell

class DirichletBocoDataset(Dataset):
    "Receives a dict of inputs and outputs and returns every possible combination of the elements in inputs \
    that must match the outputs"
    def __init__(self, inputs, outputs):
        self.inputs = np.array([inputs[k] for k in inputs])
        self.outputs = np.array([outputs[k] for k in outputs])

        # 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):
        item1, item2 = np.zeros(len(self.inputs)), np.zeros(len(self.outputs))
        for i, input in enumerate(self.inputs):
            item1[i] = input[(idx // self.mods[i]) % len(input)]
        for i, output in enumerate(self.outputs):
            item2[i] = output[(idx // self.mods[i]) % len(output)]
        return torch.from_numpy(item1).float(), torch.from_numpy(item2).float()

class DirichletBoco(Boco):
    "Dirichlet boundary condition: Fixed outputs."
    def __init__(self, name, inputs, outputs):
        super().__init__(name)
        self.type = 'dirichlet'

        # check for dict with numpy arrays
        checkValidDict(inputs)
        checkValidDict(outputs)

        # ckeck outputs has correct length
        s = 1
        for k in inputs:
            s *= len(inputs[k])
        for k in outputs:
            if len(outputs[k]) != s:
                raise Exception(f"length of {k} incorrect")


        self.inputs = inputs
        self.outputs = outputs


    def addBoco(self, input_keys, output_keys):
        # check that all inputs and outputs are present
        checkDictArray(self.inputs, input_keys)
        checkDictArray(self.outputs, output_keys)
        # reorder to match order in pde
        self.inputs = {k: self.inputs[k] for k in input_keys}
        self.outputs = {k: self.outputs[k] for k in output_keys}

    def summary(self, input_keys, output_keys, param_keys):
        super().summary()
        print('Inputs: ', self.inputs)
        print('Outputs: ', self.outputs)
        print('')

    def initialize(self):
        self.dataset = DirichletBocoDataset(self.inputs, self.outputs)
        self.dataloader = DataLoader(self.dataset, batch_size=len(self.dataset), num_workers=4)
        self.loss = torch.nn.MSELoss()
        self._inputs, self._outputs = next(iter(self.dataloader))

    def computeLoss(self, model, device):
        self._inputs, self._outputs = self._inputs.to(device), self._outputs.to(device)
        preds = model(self._inputs)
        return self.loss(preds, self._outputs)

# Cell

class NeumannBocoDataset(Dataset):
    "Receives a dict of inputs and returns every possible combination of the elements"
    def __init__(self, inputs):
        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):
        item = np.zeros(len(self.inputs))
        for i, input in enumerate(self.inputs):
            item[i] = input[(idx // self.mods[i]) % len(input)]
        return torch.from_numpy(item).float()

class NeumannBoco(Boco):
    "Neumann boundary condition: Enforces the value of certain grads to be 0"
    def __init__(self, name, inputs, grads):
        super().__init__(name)
        self.type = 'neumann'
        self.grads = grads

        # check for dict with numpy arrays
        checkValidDict(inputs)

        self.inputs = inputs

    def addBoco(self, input_keys, output_keys):
        # check that all inputs and outputs are present
        checkDictArray(self.inputs, input_keys)
        # reorder to match order in pde
        self.inputs = {k: self.inputs[k] for k in input_keys}
        self.input_keys = input_keys
        self.output_keys = output_keys

    def summary(self, input_keys, output_keys, param_keys):
        super().summary()
        print('Inputs: ', self.inputs)
        print('Grads: ', self.grads)
        print('')

    def initialize(self):
        self.dataset = NeumannBocoDataset(self.inputs)
        self.dataloader = DataLoader(self.dataset, batch_size=len(self.dataset), num_workers=4)
        self.loss = torch.nn.MSELoss()
        self._inputs = next(iter(self.dataloader))

    def computeLoss(self, model, device):
        self._inputs = self._inputs.to(device)
        self._inputs.requires_grad = True
        preds = model(self._inputs)
        # compute gradients
        _grads, = torch.autograd.grad(preds, self._inputs,
                    grad_outputs=preds.data.new(preds.shape).fill_(1),
                    create_graph=True, only_inputs=True)
        # assign keys to gradients
        grads = {
            output: {
                k: _grads[:,j] for j, k in enumerate(self.input_keys)
            } for output in self.output_keys
        }
        # compute loss for corresponding gradients
        loss = torch.cat([grads[g][self.grads[g]] for g in self.grads])
        return self.loss(loss, torch.zeros(loss.shape).to(device))