from torch import nn
from neuralpredictors.regularizers import LaplaceL2norm
import torch
import numpy as np
from utils.global_functions import model_seed, set_random_seed
from training.regularizers import l1_regularization

import torch.nn.functional as F


class FactorizedModel(nn.Module):
    def __init__(
            self,
            input_shape,
            num_of_neurons,
            num_of_channels,
            l1,
            gamma,
            do_nonlin=True,
            seed=None,
            do_sta=0,
            do_svd=False,
            svd=None,
            fancy_nonlin=False,
            sta=None,
            learn_filter=True,
            fit_gaussian=False,
            dog=False,
            sign=None,
            stride=1
    ):
        super().__init__()
        self.fit_gaussian = fit_gaussian
        self.dog = dog

        self.conv1 = nn.Conv3d(1, 1, kernel_size=(1,) + input_shape, stride=stride)
        self.conv2 = nn.Conv3d(1, num_of_neurons, kernel_size=(num_of_channels, 1, 1))
        self.init_conv(self.conv1)

        self.laplace = LaplaceL2norm()
        self.l1 = l1
        self.gamma = gamma
        self.do_sta = do_sta
        self.do_svd = do_svd
        self.svd = svd
        self.sta = sta

        if self.do_sta > 0:
            if self.do_svd:
                init_with_sta(
                    np.expand_dims(self.svd[0], 0),
                    layer=self.conv1,
                    learnable=learn_filter,
                )
                init_with_sta(
                    np.expand_dims(np.expand_dims(self.svd[1], -1), -1),
                    layer=self.conv2,
                    learnable=learn_filter,
                )
            else:
                if self.do_sta > 1:
                    init_with_sta(
                        np.expand_dims(self.sta[0], 0),
                        layer=self.conv1,
                        learnable=learn_filter,
                    )

                init_with_sta(
                    np.expand_dims(np.expand_dims(self.sta[1], -1), -1),
                    layer=self.conv2,
                    learnable=learn_filter,
                )

        self.do_nonlin = do_nonlin
        if fancy_nonlin:
            self.nonlinearity = parametrized_softplus()
        else:
            self.nonlinearity = F.softplus
        if seed is None:
            self.seed = model_seed
        else:
            self.seed = seed
        set_random_seed(self.seed)

    @staticmethod
    def init_conv(m):
        if isinstance(m, nn.Conv2d):
            nn.init.xavier_normal_(m.weight.data, gain=0.005)
            if m.bias is not None:
                m.bias.data.fill_(0)

    def forward(self, x, data_key=None):
        x = self.conv1(x)
        x = self.conv2(x)
        if self.do_nonlin:
            x = self.nonlinearity(x)
        x = torch.squeeze(x, 1).squeeze(-1)
        x = torch.flatten(x, start_dim=0, end_dim=1)
        return x

    def regularizer(self, data_key=None):
        if not self.fit_gaussian:
            return (
                self.gamma * self.laplace(self.conv1.weight[:, :, 0]),
                torch.tensor([0], device="cuda"),
                self.l1 * l1_regularization(self),
            )
        else:
            return (torch.tensor([0], device="cuda"),
                    torch.tensor([0], device="cuda"),
                    torch.tensor([0], device="cuda"))


def init_with_sta(sta, layer, learnable=False, bias=0):
    layer.weight.data = torch.tensor(sta).unsqueeze(0).unsqueeze(0)
    if layer.bias is not None:
        layer.bias.data.fill_(bias)
        layer.bias.requires_grad = learnable
    layer.weight.requires_grad = learnable

class RankModel(nn.Module):
    def __init__(
            self,
            input_shape,
            num_of_neurons,
            num_of_channels,
            l1,
            gamma,
            do_nonlin=True,
            seed=None,
            do_sta=0,
            do_svd=False,
            svd=None,
            fancy_nonlin=False,
            sta=None,
            learn_filter=True,
            fit_gaussian=False,
            dog=False,
            sign=None,
            stride=1,
            rank=2
    ):
        super().__init__()
        self.fit_gaussian = fit_gaussian
        self.dog = dog
        self.rank = rank
        self.gamma = gamma
        self.l1 = l1
        self.laplace = LaplaceL2norm()

        self.conv1 = nn.Conv3d(1, rank, kernel_size=(1,) + input_shape, stride=stride)
        self.conv2 = nn.Conv3d(rank, num_of_neurons, kernel_size=(num_of_channels, 1, 1), groups=num_of_neurons)

        self.do_nonlin = do_nonlin
        if fancy_nonlin:
            self.nonlinearity = parametrized_softplus()
        else:
            self.nonlinearity = F.softplus
        if seed is None:
            self.seed = model_seed
        else:
            self.seed = seed
        set_random_seed(self.seed)
        self.init_conv(self.conv1)
        self.init_conv(self.conv2)

    def forward(self, x, data_key=None):
        x = self.conv1(x)
        x = self.conv2(x)
        if self.do_nonlin:
            x = self.nonlinearity(x)
        x = torch.squeeze(x, 1).squeeze(-1)
        x = torch.flatten(x, start_dim=0, end_dim=1)
        return x

    @staticmethod
    def init_conv(m):
        if isinstance(m, nn.Conv3d):
            if m.bias is not None:
                m.bias.data.fill_(0)

    def regularizer(self, data_key=None):
        if not self.fit_gaussian:
            return (
                self.gamma * self.laplace(self.conv1.weight[:, :, 0]),
                torch.tensor([0], device="cuda"),
                self.l1 * l1_regularization(self),
            )
        else:
            return (torch.tensor([0], device="cuda"),
                    torch.tensor([0], device="cuda"),
                    torch.tensor([0], device="cuda"))



class LearnedSoftPlus(torch.nn.Module):
    def __init__(self, init_beta=1.0, threshold=20):
        super().__init__()
        # keep beta > 0
        self.log_beta = torch.nn.Parameter(torch.tensor(float(init_beta)).log())
        self.threshold = 20

    def forward(self, x):
        beta = self.log_beta.exp()
        beta_x = beta * x
        return torch.where(beta_x < 20, torch.log1p(beta_x.exp()) / beta, x)


class ParametrizedELU(nn.Module):
    def __init__(self, a=1.0, b=1.0, c=1.0, device='cuda', *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.a = torch.nn.Parameter(torch.tensor(a, device=device), requires_grad=True)
        self.b = torch.nn.Parameter(torch.tensor(b, device=device), requires_grad=True)
        self.c = torch.nn.Parameter(torch.tensor(c, device=device), requires_grad=True)

    def forward(self, x):
        return torch.where(x > 0, self.c * x, self.a * (torch.exp(x / self.b) - 1))


class paremetrized_exp(nn.Module):
    """
    Implementation of parametrized exponential activation.
    """

    def __init__(self, alpha=None):
        super().__init__()
        self.weight_clamg = ParamClamp()
        if alpha is None:
            self.alpha = nn.Parameter(torch.tensor(1.0), requires_grad=True)
        else:
            self.alpha = nn.Parameter(torch.tensor(alpha), requires_grad=True)

        self.weight_clamg.apply(self.alpha)

    def forward(self, x):
        return torch.exp(self.alpha * x)


class ParamClamp(torch.autograd.Function):
    def forward(ctx, input, **kwargs):
        return input.clamp(min=0)

    def backward(ctx, grad_output):
        return grad_output


class parametrized_elu(nn.Module):
    def __init__(self, a=1.0, b=1.0, w=1.0):
        super().__init__()
        self.a = nn.Parameter(torch.tensor(a), requires_grad=True)
        self.b = nn.Parameter(torch.tensor(b), requires_grad=True)
        self.w = nn.Parameter(torch.tensor(w), requires_grad=True)

    def forward(self, x):
        return torch.where(x >= 0, self.w * x, self.a * (torch.exp(x / self.b) - 1))


class parametrized_softplus(nn.Module):
    def __init__(self, a=1.0, b=0.0, w=1.0):
        super().__init__()
        self.a = nn.Parameter(torch.tensor(a), requires_grad=True)
        self.b = nn.Parameter(torch.tensor(b), requires_grad=True)
        self.w = nn.Parameter(torch.tensor(w), requires_grad=True)

    def forward(self, x):
        return self.a * torch.log(1 + torch.exp(self.w * x + self.b))


class multi_softplus(nn.ModuleDict):
    def __init__(self, num_of_cells, data_keys=None):
        super().__init__()
        self.a, self.b, self.w = nn.ParameterDict(), nn.ParameterDict(), nn.ParameterDict()
        for data_key in data_keys:
            self.a[data_key] = nn.Parameter(torch.ones(num_of_cells[data_key]), requires_grad=True)
            self.b[data_key] = nn.Parameter(torch.zeros(num_of_cells[data_key]), requires_grad=True)
            self.w[data_key] = nn.Parameter(torch.ones(num_of_cells[data_key]), requires_grad=True)
        self.data_keys = data_keys

    def forward(self, x, data_key):
        if data_key is None:
            data_key = list(self.data_keys)[0]
        a_expanded = self.a[data_key].unsqueeze(0).expand_as(x)
        b_expanded = self.b[data_key].unsqueeze(0).expand_as(x)
        w_expanded = self.w[data_key].unsqueeze(0).expand_as(x)

        # Apply custom parametrized softplus
        output = a_expanded * torch.log(1 + torch.exp(w_expanded * x + b_expanded))

        return output


