import torch

from pina.problem import AbstractProblem
from pina import Condition, LabelTensor
from pina.solvers import GAROM
from pina.trainer import Trainer
import torch.nn as nn
import matplotlib.tri as tri


def func(x, mu1, mu2):
    import torch
    x_m1 = (x[:, 0] - mu1).pow(2)
    x_m2 = (x[:, 1] - mu2).pow(2)
    norm = x[:, 0]**2 + x[:, 1]**2
    return torch.exp(-(x_m1 + x_m2))


class ParametricGaussian(AbstractProblem):
    output_variables = [f'u_{i}' for i in range(900)]

    # params
    xx = torch.linspace(-1, 1, 20)
    yy = xx
    params = LabelTensor(torch.cartesian_prod(xx, yy), labels=['mu1', 'mu2'])

    # define domain
    x = torch.linspace(-1, 1, 30)
    domain = torch.cartesian_prod(x, x)
    triang = tri.Triangulation(domain[:, 0], domain[:, 1])
    sol = []
    for p in params:
        sol.append(func(domain, p[0], p[1]))
    snapshots = LabelTensor(torch.stack(sol), labels=output_variables)

    # define conditions
    conditions = {
        'data': Condition(input_points=params, output_points=snapshots)
    }


# simple Generator Network
class Generator(nn.Module):

    def __init__(self,
                 input_dimension,
                 parameters_dimension,
                 noise_dimension,
                 activation=torch.nn.SiLU):
        super().__init__()

        self._noise_dimension = noise_dimension
        self._activation = activation

        self.model = torch.nn.Sequential(
            torch.nn.Linear(6 * self._noise_dimension, input_dimension // 6),
            self._activation(),
            torch.nn.Linear(input_dimension // 6, input_dimension // 3),
            self._activation(),
            torch.nn.Linear(input_dimension // 3, input_dimension))
        self.condition = torch.nn.Sequential(
            torch.nn.Linear(parameters_dimension, 2 * self._noise_dimension),
            self._activation(),
            torch.nn.Linear(2 * self._noise_dimension,
                            5 * self._noise_dimension))

    def forward(self, param):
        # uniform sampling in [-1, 1]
        z = torch.rand(size=(param.shape[0], self._noise_dimension),
                       device=param.device,
                       dtype=param.dtype,
                       requires_grad=True)
        z = 2. * z - 1.

        # conditioning by concatenation of mapped parameters
        input_ = torch.cat((z, self.condition(param)), dim=-1)
        out = self.model(input_)

        return out


# Simple Discriminator Network
class Discriminator(nn.Module):

    def __init__(self,
                 input_dimension,
                 parameter_dimension,
                 hidden_dimension,
                 activation=torch.nn.ReLU):
        super().__init__()

        self._activation = activation
        self.encoding = torch.nn.Sequential(
            torch.nn.Linear(input_dimension, input_dimension // 3),
            self._activation(),
            torch.nn.Linear(input_dimension // 3, input_dimension // 6),
            self._activation(),
            torch.nn.Linear(input_dimension // 6, hidden_dimension))
        self.decoding = torch.nn.Sequential(
            torch.nn.Linear(2 * hidden_dimension, input_dimension // 6),
            self._activation(),
            torch.nn.Linear(input_dimension // 6, input_dimension // 3),
            self._activation(),
            torch.nn.Linear(input_dimension // 3, input_dimension),
        )

        self.condition = torch.nn.Sequential(
            torch.nn.Linear(parameter_dimension, hidden_dimension // 2),
            self._activation(),
            torch.nn.Linear(hidden_dimension // 2, hidden_dimension))

    def forward(self, data):
        x, condition = data
        encoding = self.encoding(x)
        conditioning = torch.cat((encoding, self.condition(condition)), dim=-1)
        decoding = self.decoding(conditioning)
        return decoding


problem = ParametricGaussian()


def test_constructor():
    GAROM(problem=problem,
          generator=Generator(input_dimension=900,
                              parameters_dimension=2,
                              noise_dimension=12),
          discriminator=Discriminator(input_dimension=900,
                                      parameter_dimension=2,
                                      hidden_dimension=64))


def test_train_cpu():
    solver = GAROM(problem=problem,
                   generator=Generator(input_dimension=900,
                                       parameters_dimension=2,
                                       noise_dimension=12),
                   discriminator=Discriminator(input_dimension=900,
                                               parameter_dimension=2,
                                               hidden_dimension=64))

    trainer = Trainer(solver=solver, max_epochs=4, accelerator='cpu', batch_size=20)
    trainer.train()


def test_sample():
    solver = GAROM(problem=problem,
                   generator=Generator(input_dimension=900,
                                       parameters_dimension=2,
                                       noise_dimension=12),
                   discriminator=Discriminator(input_dimension=900,
                                               parameter_dimension=2,
                                               hidden_dimension=64))
    solver.sample(problem.params)
    assert solver.sample(problem.params).shape == problem.snapshots.shape


def test_forward():
    solver = GAROM(problem=problem,
                   generator=Generator(input_dimension=900,
                                       parameters_dimension=2,
                                       noise_dimension=12),
                   discriminator=Discriminator(input_dimension=900,
                                               parameter_dimension=2,
                                               hidden_dimension=64))
    solver(problem.params, mc_steps=100, variance=True)
    assert solver(problem.params).shape == problem.snapshots.shape