import math

import torch
import torch.nn.functional as F
from fastai.layers import *
from fastcore.basics import snake2camel
from fastcore.meta import delegates
from torch.nn.init import normal_
from torch.nn.utils import spectral_norm, weight_norm

from fedot_ind.core.architecture.settings.computational import backend_methods as np
from fedot_ind.core.models.nn.network_modules.layers.linear_layers import AddCoords1d
from fedot_ind.core.models.nn.network_modules.layers.padding_layers import *


class Conv2dSame(Module):
    """
    Conv2d with padding='same'
    """

    def __init__(self, ni, nf, ks=(3, 3), stride=(1, 1), dilation=(1, 1), **kwargs):
        if isinstance(ks, Integral):
            ks = (ks, ks)
        if isinstance(stride, Integral):
            stride = (stride, stride)
        if isinstance(dilation, Integral):
            dilation = (dilation, dilation)
        self.ks, self.stride, self.dilation = ks, stride, dilation
        self.conv2d_same = nn.Conv2d(
            ni, nf, ks, stride=stride, dilation=dilation, **kwargs)
        self.weight = self.conv2d_same.weight
        self.bias = self.conv2d_same.bias
        self.pad = Pad2d

    def forward(self, x):
        self.padding = same_padding2d(x.shape[-2], x.shape[-1], self.ks,
                                      dilation=self.dilation)  # stride=self.stride not used in padding calculation!
        return self.conv2d_same(self.pad(self.padding)(x))


@delegates(nn.Conv2d.__init__)
def Conv2d(ni, nf, kernel_size=None, ks=None, stride=1, padding='same', dilation=1, init='auto', bias_std=0.01,
           **kwargs):
    """conv1d layer with padding='same', 'valid', or any integer (defaults to 'same')"""
    assert not (
        kernel_size and ks), 'use kernel_size or ks but not both simultaneously'
    assert kernel_size is not None or ks is not None, 'you need to pass a ks'
    kernel_size = kernel_size or ks
    if padding == 'same':
        conv = Conv2dSame(ni, nf, kernel_size, stride=stride,
                          dilation=dilation, **kwargs)
    elif padding == 'valid':
        conv = nn.Conv2d(ni, nf, kernel_size, stride=stride,
                         padding=0, dilation=dilation, **kwargs)
    else:
        conv = nn.Conv2d(ni, nf, kernel_size, stride=stride,
                         padding=padding, dilation=dilation, **kwargs)
    init_linear(conv, None, init=init, bias_std=bias_std)
    return conv


class CausalConv1d(torch.nn.Conv1d):
    def __init__(self, ni, nf, ks, stride=1, dilation=1, groups=1, bias=True):
        super(CausalConv1d, self).__init__(ni, nf, kernel_size=ks, stride=stride, padding=0, dilation=dilation,
                                           groups=groups, bias=bias)
        self.__padding = (ks - 1) * dilation

    def forward(self, input):
        return super(CausalConv1d, self).forward(F.pad(input, (self.__padding, 0)))


@delegates(nn.Conv1d.__init__)
def Conv1d(ni, nf, kernel_size=None, ks=None, stride=1, padding='same', dilation=1, init='auto', bias_std=0.01,
           **kwargs):
    """conv1d layer with padding='same', 'causal', 'valid', or any integer (defaults to 'same')"""
    assert not (
        kernel_size and ks), 'use kernel_size or ks but not both simultaneously'
    assert kernel_size is not None or ks is not None, 'you need to pass a ks'
    kernel_size = kernel_size or ks
    if padding == 'same':
        if kernel_size % 2 == 1:
            conv = nn.Conv1d(ni, nf, kernel_size, stride=stride, padding=kernel_size // 2 * dilation, dilation=dilation,
                             **kwargs)
        else:
            conv = SameConv1d(ni, nf, kernel_size,
                              stride=stride, dilation=dilation, **kwargs)
    elif padding == 'causal':
        conv = CausalConv1d(ni, nf, kernel_size,
                            stride=stride, dilation=dilation, **kwargs)
    elif padding == 'valid':
        conv = nn.Conv1d(ni, nf, kernel_size, stride=stride,
                         padding=0, dilation=dilation, **kwargs)
    else:
        conv = nn.Conv1d(ni, nf, kernel_size, stride=stride,
                         padding=padding, dilation=dilation, **kwargs)
    init_linear(conv, None, init=init, bias_std=bias_std)
    return conv


class SeparableConv1d(Module):
    def __init__(self, ni, nf, ks, stride=1, padding='same', dilation=1, bias=True, bias_std=0.01):
        self.depthwise_conv = Conv1d(ni, ni, ks, stride=stride, padding=padding, dilation=dilation, groups=ni,
                                     bias=bias)
        self.pointwise_conv = nn.Conv1d(
            ni, nf, 1, stride=1, padding=0, dilation=1, groups=1, bias=bias)
        if bias:
            if bias_std != 0:
                normal_(self.depthwise_conv.bias, 0, bias_std)
                normal_(self.pointwise_conv.bias, 0, bias_std)
            else:
                self.depthwise_conv.bias.data.zero_()
                self.pointwise_conv.bias.data.zero_()

    def forward(self, x):
        x = self.depthwise_conv(x)
        x = self.pointwise_conv(x)
        return x


def SEModule1d(ni, reduction=16, act=nn.ReLU, act_kwargs={}):
    """Squeeze and excitation module for 1d"""
    nf = math.ceil(ni // reduction / 8) * 8
    assert nf != 0, 'nf cannot be 0'
    return SequentialEx(nn.AdaptiveAvgPool1d(1),
                        ConvBlock(ni, nf, ks=1, norm=None,
                                  act=act, act_kwargs=act_kwargs),
                        ConvBlock(nf, ni, ks=1, norm=None, act=nn.Sigmoid), ProdLayer())


class ConvBlock(nn.Sequential):
    """Create a sequence of conv1d (`ni` to `nf`), activation (if `act_cls`) and `norm_type` layers."""

    def __init__(self, ni, nf, kernel_size=None, ks=3, stride=1, padding='same', bias=None, bias_std=0.01, norm='Batch',
                 zero_norm=False, batch_norm_1st=True,
                 act=nn.ReLU, act_kwargs={}, init='auto', dropout=0., xtra=None, coord=False, separable=False,
                 **kwargs):
        kernel_size = kernel_size or ks
        ndim = 1
        layers = [AddCoords1d()] if coord else []
        norm_type = getattr(
            NormType, f"{snake2camel(norm)}{'Zero' if zero_norm else ''}") if norm is not None else None
        batch_norm = norm_type in (NormType.Batch, NormType.BatchZero)
        inn = norm_type in (NormType.Instance, NormType.InstanceZero)
        if bias is None:
            bias = not (batch_norm or inn)
        if separable:
            conv = SeparableConv1d(ni + coord, nf, ks=kernel_size,
                                   bias=bias, stride=stride, padding=padding, **kwargs)
        else:
            conv = Conv1d(ni + coord, nf, ks=kernel_size, bias=bias,
                          stride=stride, padding=padding, **kwargs)
        act = None if act is None else act(**act_kwargs)
        if not separable:
            init_linear(conv, act, init=init, bias_std=bias_std)
        if norm_type == NormType.Weight:
            conv = weight_norm(conv)
        elif norm_type == NormType.Spectral:
            conv = spectral_norm(conv)
        layers += [conv]
        act_batch_norm = []
        if act is not None:
            act_batch_norm.append(act)
        if batch_norm:
            act_batch_norm.append(
                BatchNorm(nf, norm_type=norm_type, ndim=ndim))
        if inn:
            act_batch_norm.append(InstanceNorm(
                nf, norm_type=norm_type, ndim=ndim))
        if batch_norm_1st:
            act_batch_norm.reverse()
        if dropout:
            layers += [nn.Dropout(dropout)]
        layers += act_batch_norm
        if xtra:
            layers.append(xtra)
        super().__init__(*layers)


class MultiConv1d(Module):
    """Module that applies multiple convolutions with different kernel sizes"""

    def __init__(self, ni, nf=None, kss=[1, 3, 5, 7], keep_original=False, separable=False, dim=1, **kwargs):
        kss = list(kss)
        n_layers = len(kss)
        if ni == nf:
            keep_original = False
        if nf is None:
            nf = ni * (keep_original + n_layers)
        nfs = [(nf - ni * keep_original) // n_layers] * n_layers
        while np.sum(nfs) + ni * keep_original < nf:
            for i in range(len(nfs)):
                nfs[i] += 1
                if np.sum(nfs) + ni * keep_original == nf:
                    break

        _conv = SeparableConv1d if separable else Conv1d
        self.layers = nn.ModuleList()
        for nfi, ksi in zip(nfs, kss):
            self.layers.append(_conv(ni, nfi, ksi, **kwargs))
        self.keep_original, self.dim = keep_original, dim

    def forward(self, x):
        output = [x] if self.keep_original else []
        for l in self.layers:
            output.append(l(x))
        x = torch.cat(output, dim=self.dim)
        return x
