import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from copy import deepcopy
import utils
from einops import rearrange, repeat
from collections import OrderedDict
from probconserv import apply_constraint, get_convex_params
# from tqdm.notebook import tqdm

# This code is borrowed from FNO git repository: https://github.com/zongyi-li/fourier_neural_operator


device = "cuda" if torch.cuda.is_available() else "cpu" 

################################################################
# 2D fourier layer
################################################################
class SpectralConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, modes1, modes2, drop_modes=0., drop_channels=0.):
        super(SpectralConv2d, self).__init__()

        """
        2D Fourier layer. It does FFT, linear transform, and Inverse FFT.    
        """

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.modes1 = modes1  # Number of Fourier modes to multiply, at most floor(N/2) + 1
        self.modes2 = modes2

        self.scale = (1 / (in_channels * out_channels))
        self.weights1 = nn.Parameter(
            self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))
        self.weights2 = nn.Parameter(
            self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))
        
        self.binomial_modes = torch.distributions.Binomial(self.modes1, probs=torch.tensor(1-drop_modes))

        self.n_drop_channels = int(drop_channels * self.in_channels)

    # Complex multiplication
    def compl_mul2d(self, input, weights):
        # (batch, in_channel, x,y ), (in_channel, out_channel, x,y) -> (batch, out_channel, x,y)
        return torch.einsum("bixy,ioxy->boxy", input, weights)

    def forward(self, x):
        batchsize = x.shape[0]
        # Compute Fourier coeffcients up to factor of e^(- something constant)
        x_ft = torch.fft.rfft2(x)

        n_active_modes = self.binomial_modes.sample().int()
        channel_mask = torch.ones_like(self.weights1)
        perm = torch.randperm(self.in_channels)
        idx = perm[:self.n_drop_channels]
        channel_mask[idx, :, :] = 0.
        new_weights1 = self.weights1 * channel_mask
        new_weights2 = self.weights2 * channel_mask
        # n_active_modes = self.modes1
        # new_weights1 = self.weights1
        # new_weights2 = self.weights2

        # Multiply relevant Fourier modes
        out_ft = torch.zeros(batchsize, self.out_channels, x.size(-2), x.size(-1) // 2 + 1, dtype=torch.cfloat,
                             device=x.device)
        out_ft[:, :, :n_active_modes, :self.modes2] = \
            self.compl_mul2d(x_ft[:, :, :n_active_modes, :self.modes2], new_weights1[:, :, :n_active_modes])
        out_ft[:, :, -n_active_modes:, :self.modes2] = \
            self.compl_mul2d(x_ft[:, :, -n_active_modes:, :self.modes2], new_weights2[:, :, :n_active_modes])

        # Return to physical space
        x = torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1)))
        return x


class FourierBlock(nn.Module):
    def __init__(self, modes1, modes2, width, drop_modes=0., drop_channels=0., activation=True):
        super().__init__()

        self.conv = SpectralConv2d(width, width, modes1, modes2, drop_modes, drop_channels)
        self.w = nn.Conv2d(width, width, 1)
        self.activation = activation

    def forward(self, x):
        x1 = self.conv(x)
        x2 = self.w(x)
        x = x1 + x2
        if self.activation:
            x = F.gelu(x)
        return x


class FNO2d(nn.Module):
    def __init__(self,
                 modes1,
                 modes2,
                 width,
                 n_layers=4,
                 dropout=0,
                 drop_modes=0.,
                 drop_channels=0.,
                 output_var=False,
                 cov_type="diagonal",  # "diagonal", "lowrank", "full"
                 rank=4,
                 n_outputs=1,
                 last_layer_reshape=False,
                 lb=0,
                 ub=1,
                 probconserv=False,
                 probconserv_sampling=False,
                 nonlinear_constraint="none",
                 t=None,
                 tpred=None,
                 grid_train=None,
                 dataset_class=None):
        super(FNO2d, self).__init__()

        self.modes1 = modes1
        self.modes2 = modes2
        self.width = width
        self.lb = lb
        self.ub = ub
        self.dropout = dropout
        self.output_var = output_var
        self.n_outputs = n_outputs
        self.rank = rank
        self.probconserv = probconserv
        self.probconserv_sampling = probconserv_sampling
        self.nonlinear_constraint = nonlinear_constraint
        self.cov_type = cov_type
        
        self.t = t
        self.tpred = tpred
        self.grid_train = grid_train
        self.dataset_class = dataset_class

        self.padding = 9
        self.fc0 = nn.Linear(3, self.width)

        self.fourier_layers = nn.Sequential(
            *[FourierBlock(self.modes1, self.modes2, self.width, drop_modes, drop_channels) for _ in range(n_layers - 1)],
            FourierBlock(self.modes1, self.modes2, self.width, drop_modes, drop_channels, activation=False)
        )

        self.fc1 = nn.Linear(self.width, 128)

        if self.output_var:
            if self.cov_type == "diagonal":
                out_dim = 2 * n_outputs
            elif self.cov_type == "lowrank":
                out_dim = n_outputs + n_outputs * self.rank
            elif self.cov_type == "full":
                out_dim = n_outputs + (n_outputs * (n_outputs + 1)) // 2
            else:
                raise ValueError("Unsupported covariance type")
        else:
            out_dim = n_outputs

        self.fc2 = nn.Linear(128, out_dim)

        if self.output_var:
            if self.cov_type == "diagonal":
                self.loss_func = utils.nll_mu_var
                self.loss_func = utils.CRPSLoss()
            else:
                self.loss_func = utils.nll_mu_cov
        else:
            self.loss_func = utils.LpLoss(size_average=False)

        self.last_layer_reshape = last_layer_reshape

    def _apply_constraints(self, mu, std, x, t, tpred, grid_train, dataset_class):
        if self.cov_type != "diagonal":
            raise NotImplementedError("Constraint logic currently supports only diagonal covariance.")

        if self.probconserv:
            if self.probconserv_sampling:
                n_samples = 100
                nf, nx, nt, _ = mu.shape

                dist = torch.distributions.Normal(mu, std)
                samples = dist.rsample((n_samples,))  # shape: [nb, nf, nx, nt, 1]

                samples = rearrange(samples, "nb nf nx nt 1 -> (nb nf) nx nt 1")
                rep_std = repeat(std, "nf nx nt 1 -> (nb nf) nx nt 1", nb=n_samples)
                rep_x = repeat(x, "nf nx nt 1 -> (nb nf) nx nt 1", nb=n_samples)

                mass_rhs_func = dataset_class.get_mass_rhs_func(x=rep_x)

                new_samples, _, _, _ = apply_constraint(
                    mu=samples[..., 0],
                    std=rep_std[..., 0],
                    mass_rhs_func=mass_rhs_func,
                    t=t,
                    tpred=tpred,
                    grid_train=grid_train,
                    precis_g=np.inf,
                    second_deriv_alpha=None
                )

                new_samples = rearrange(new_samples, "(nb nf) nx nt -> nb nf nx nt", nb=n_samples)
                sample_mu = new_samples.mean(dim=0)
                sample_std = new_samples.std(dim=0)

                mu = sample_mu
                std = sample_std

            else:
                mass_rhs_func = dataset_class.get_mass_rhs_func(x=x)
                mu, std, _, _ = apply_constraint(
                    mu=mu[..., 0],
                    std=std[..., 0],
                    mass_rhs_func=mass_rhs_func,
                    t=t,
                    tpred=tpred,
                    grid_train=grid_train,
                    precis_g=np.inf,
                    second_deriv_alpha=None
                )
        else:
            mu = mu[..., 0]
            std = std[..., 0]


        if self.nonlinear_constraint == "none":
            return mu.unsqueeze(-1), std.square().unsqueeze(-1)

        nf, nx, nt = mu.shape
        mu = mu.reshape(nf, nx * nt)
        std = std.reshape(nf, nx * nt)

        if self.nonlinear_constraint == "sampling":
            num_samples = 5
            dist = torch.distributions.Normal(mu, std)
            y_samples = dist.rsample((num_samples,))
            y_samples = y_samples / torch.norm(y_samples, dim=-1, keepdim=True)
            mu = y_samples.mean(dim=0)
            std = y_samples.std(dim=0)

        elif self.nonlinear_constraint == "deltamethod":
            norm_mu = torch.norm(mu, dim=1, keepdim=True)
            mu = mu / norm_mu
            I = torch.eye(mu.shape[1], device=mu.device).expand(mu.shape[0], -1, -1)
            J_g = I / norm_mu.unsqueeze(-1) - (mu.unsqueeze(-1) @ mu.unsqueeze(1)) / norm_mu.unsqueeze(-1).pow(3)
            Sigma_diag = std ** 2
            J_Sigma = J_g ** 2 * Sigma_diag.unsqueeze(-2)
            var_diag = torch.sum(J_Sigma, dim=-1)
            std = torch.sqrt(torch.clamp(var_diag, min=1e-6))

        mu = mu.view(nf, nx, nt)
        std = std.view(nf, nx, nt)
        return mu.unsqueeze(-1), std.square().unsqueeze(-1)

    def forward(self, x):
        grid = self.get_grid(x.shape, x.device)
        x = torch.cat((x, grid), dim=-1)
        x = self.fc0(x)
        x = x.permute(0, 3, 1, 2)
        x = F.pad(x, [0, self.padding, 0, self.padding])

        x = self.fourier_layers(x)
        x = x[..., :-self.padding, :-self.padding]
        x = x.permute(0, 2, 3, 1)
        x = self.fc1(x)
        x = F.dropout(x, training=True, p=self.dropout)
        x = F.gelu(x)

        if self.last_layer_reshape:
            x = rearrange(x, 'b p t d -> (b p t) d')

        x = self.fc2(x)

        if self.output_var:
            if self.cov_type == "diagonal":
                mu, logvar = torch.chunk(x, 2, dim=-1)
                var = F.softplus(logvar)
                return mu, var
            elif self.cov_type == "lowrank":
                mu = x[..., :self.n_outputs]
                L = x[..., self.n_outputs:].view(*x.shape[:-1], self.n_outputs, self.rank)
                cov = torch.matmul(L, L.transpose(-1, -2))
                return mu, cov
            elif self.cov_type == "full":
                mu = x[..., :self.n_outputs]
                tril_indices = torch.tril_indices(row=self.n_outputs, col=self.n_outputs, offset=0)
                L_elements = x[..., self.n_outputs:]
                L = torch.zeros(*x.shape[:-1], self.n_outputs, self.n_outputs, device=x.device)
                L[..., tril_indices[0], tril_indices[1]] = L_elements
                cov = torch.matmul(L, L.transpose(-1, -2))
                return mu, cov

        return x

    def get_grid(self, shape, device):
        batchsize, size_x, size_y = shape[0], shape[1], shape[2]
        gridx = torch.tensor(np.linspace(self.lb, self.ub, size_x), dtype=torch.float)
        gridx = gridx.reshape(1, size_x, 1, 1).repeat([batchsize, 1, size_y, 1])
        gridy = torch.tensor(np.linspace(self.lb, self.ub, size_y), dtype=torch.float)
        gridy = gridy.reshape(1, 1, size_y, 1).repeat([batchsize, size_x, 1, 1])
        return torch.cat((gridx, gridy), dim=-1).to(device)

    def test(self, test_loader, **test_params):
        self.eval()
        test_l2 = 0.0
        tpred = test_params.get("tpred", None)
        t = test_params.get("t", None)
        dataset_class = test_params.get("dataset_class", None)
        grid_train = test_params.get("grid_train", None)

        with torch.no_grad():
            for batch in test_loader:
                x, y = batch
                x, y = x.to(device), y.to(device)
                out = self(x)

                if self.probconserv or self.nonlinear_constraint != "none":
                    mu, var = out
                    std = torch.sqrt(var)
                    out = self._apply_constraints(mu, std, x, t, tpred, grid_train, dataset_class)

                test_l2 += self.loss_func(out, y).item()

        test_l2 /= len(test_loader.dataset)
        return {"loss": test_l2}

    def fit(self, train_loader, valid_loader, **fit_params):
        lr = fit_params.get("lr", 1e-3)
        step_size = fit_params.get("step_size", 50)
        gamma = fit_params.get("gamma", 0.5)
        epochs = fit_params.get("epochs", 200)
        tpred = fit_params.get("tpred", None)
        t = fit_params.get("t", None)
        dataset_class = fit_params.get("dataset_class", None)
        grid_train = fit_params.get("grid_train", None)

        optimizer = torch.optim.Adam(self.parameters(), lr=lr, weight_decay=1e-4)
        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)

        best_valid_l2 = np.inf

        for epoch in range(epochs):
            self.train()
            train_l2 = 0
            for batch in train_loader:
                x, y = batch
                x, y = x.to(device), y.to(device)
                optimizer.zero_grad()
                out = self(x)

                if self.probconserv or self.nonlinear_constraint != "none":
                    mu, var = out
                    std = torch.sqrt(var)
                    out = self._apply_constraints(mu, std, x, t, tpred, grid_train, dataset_class)

                l2 = self.loss_func(out, y) 
                #+ 1e-3/torch.mean(out[1])
                l2.backward()
                # print(out[1].grad_fn)
                optimizer.step()
                train_l2 += l2.item()

            train_l2 /= len(train_loader.dataset)
            scheduler.step()
            if valid_loader is not None:
                valid_l2 = self.test(valid_loader, **fit_params)["loss"]
            else:
                valid_l2 = train_l2

            saved = ""
            if valid_l2 < best_valid_l2:
                best_valid_l2 = valid_l2
                best_model_state_dict = deepcopy(self.state_dict())
                saved = "(saved)"

            print(f"Epoch {epoch}: Train loss={train_l2:.6f}, Validation loss={valid_l2:.6f} {saved}")

        self.load_state_dict(best_model_state_dict)
        train_l2 = self.test(train_loader, **fit_params)["loss"]
        if valid_loader is not None:
            valid_l2 = self.test(valid_loader, **fit_params)["loss"]
        else:
            valid_l2 = train_l2
        print(f"Finished training with best train loss: {train_l2:.6f} and validation loss: {valid_l2:.6f}")