import torch
import torch.nn as nn
from spikingjelly.clock_driven import neuron
from spikingjelly.clock_driven.neuron import MultiStepLIFNode
from spikingjelly.clock_driven import surrogate

from timm.models.layers import to_2tuple, trunc_normal_, DropPath
from timm.models.registry import register_model
from timm.models.vision_transformer import _cfg
import torch.nn.functional as F
from functools import partial
# __all__ = ['spikformer']


a = 1.0

class MLP(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1_linear = nn.Linear(in_features, hidden_features)
        self.fc1_bn = nn.BatchNorm1d(hidden_features)
        self.fc1_lif = MixedLIF(tau=2.0, detach_reset=True, backend='cupy')

        self.fc2_linear = nn.Linear(hidden_features, out_features)
        self.fc2_bn = nn.BatchNorm1d(out_features)
        self.fc2_lif = MixedLIF(tau=2.0, detach_reset=True, backend='cupy')

        self.c_hidden = hidden_features
        self.c_output = out_features

    def forward(self, x):
        T,B,N,C = x.shape
        x_ = x.flatten(0, 1)
        x = self.fc1_linear(x_)
        x = self.fc1_bn(x.transpose(-1, -2)).transpose(-1, -2).reshape(T, B, N, self.c_hidden).contiguous()
        x = self.fc1_lif(x)

        x = self.fc2_linear(x.flatten(0,1))
        x = self.fc2_bn(x.transpose(-1, -2)).transpose(-1, -2).reshape(T, B, N, C).contiguous()
        x = self.fc2_lif(x)
        return x


class SSA(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1):
        super().__init__()
        assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."
        self.dim = dim
        self.num_heads = num_heads
        self.scale = 0.125
        self.q_linear = nn.Linear(dim, dim)
        self.q_bn = nn.BatchNorm1d(dim)
        self.q_lif = MixedLIF(tau=2.0, detach_reset=True, backend='cupy')

        self.k_linear = nn.Linear(dim, dim)
        self.k_bn = nn.BatchNorm1d(dim)
        self.k_lif = MixedLIF(tau=2.0, detach_reset=True, backend='cupy')

        self.v_linear = nn.Linear(dim, dim)
        self.v_bn = nn.BatchNorm1d(dim)
        self.v_lif = MixedLIF(tau=2.0, detach_reset=True, backend='cupy')
        self.attn_lif = MixedLIF(tau=2.0, v_threshold=0.5, detach_reset=True, backend='cupy')

        self.proj_linear = nn.Linear(dim, dim)
        self.proj_bn = nn.BatchNorm1d(dim)
        self.proj_lif = MixedLIF(tau=2.0, detach_reset=True, backend='cupy')

    def forward(self, x):
        T,B,N,C = x.shape

        x_for_qkv = x.flatten(0, 1)  # TB, N, C
        q_linear_out = self.q_linear(x_for_qkv)  # [TB, N, C]
        q_linear_out = self.q_bn(q_linear_out. transpose(-1, -2)).transpose(-1, -2).reshape(T, B, N, C).contiguous()
        q_linear_out = self.q_lif(q_linear_out)
        q = q_linear_out.reshape(T, B, N, self.num_heads, C//self.num_heads).permute(0, 1, 3, 2, 4).contiguous()

        k_linear_out = self.k_linear(x_for_qkv)
        k_linear_out = self.k_bn(k_linear_out. transpose(-1, -2)).transpose(-1, -2).reshape(T, B, N, C).contiguous()
        k_linear_out = self.k_lif(k_linear_out)
        k = k_linear_out.reshape(T, B, N, self.num_heads, C//self.num_heads).permute(0, 1, 3, 2, 4).contiguous()

        v_linear_out = self.v_linear(x_for_qkv)
        v_linear_out = self.v_bn(v_linear_out. transpose(-1, -2)).transpose(-1, -2).reshape(T, B, N, C).contiguous()
        v_linear_out = self.v_lif(v_linear_out)
        v = v_linear_out.reshape(T, B, N, self.num_heads, C//self.num_heads).permute(0, 1, 3, 2, 4).contiguous()

        attn = (q @ k.transpose(-2, -1)) * self.scale
        x = attn @ v
        x = x.transpose(2, 3).reshape(T, B, N, C).contiguous()
        x = self.attn_lif(x)
        x = x.flatten(0, 1)
        x = self.proj_lif(self.proj_bn(self.proj_linear(x).transpose(-1, -2)).transpose(-1, -2).reshape(T, B, N, C))
        return x

class Block(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., norm_layer=nn.LayerNorm, sr_ratio=1):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = SSA(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
                              attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio)
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = MLP(in_features=dim, hidden_features=mlp_hidden_dim, drop=drop)

    def forward(self, x):
        x = x + self.attn(x)
        x = x + self.mlp(x)
        return x


class SPS(nn.Module):
    def __init__(self, img_size_h=128, img_size_w=128, patch_size=4, in_channels=2, embed_dims=256):
        super().__init__()
        self.image_size = [img_size_h, img_size_w]
        patch_size = to_2tuple(patch_size)
        self.patch_size = patch_size
        self.C = in_channels
        self.H, self.W = self.image_size[0] // patch_size[0], self.image_size[1] // patch_size[1]
        self.num_patches = self.H * self.W
        self.proj_conv = nn.Conv2d(in_channels, embed_dims//8, kernel_size=3, stride=1, padding=1, bias=False)
        self.proj_bn = nn.BatchNorm2d(embed_dims//8)
        self.proj_lif = MixedLIF(tau=2.0, detach_reset=True, backend='cupy')

        self.proj_conv1 = nn.Conv2d(embed_dims//8, embed_dims//4, kernel_size=3, stride=1, padding=1, bias=False)
        self.proj_bn1 = nn.BatchNorm2d(embed_dims//4)
        self.proj_lif1 = MixedLIF(tau=2.0, detach_reset=True, backend='cupy')

        self.proj_conv2 = nn.Conv2d(embed_dims//4, embed_dims//2, kernel_size=3, stride=1, padding=1, bias=False)
        self.proj_bn2 = nn.BatchNorm2d(embed_dims//2)
        self.proj_lif2 = MixedLIF(tau=2.0, detach_reset=True, backend='cupy')
        self.maxpool2 = torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)

        self.proj_conv3 = nn.Conv2d(embed_dims//2, embed_dims, kernel_size=3, stride=1, padding=1, bias=False)
        self.proj_bn3 = nn.BatchNorm2d(embed_dims)
        self.proj_lif3 = MixedLIF(tau=2.0, detach_reset=True, backend='cupy')
        self.maxpool3 = torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)

        self.rpe_conv = nn.Conv2d(embed_dims, embed_dims, kernel_size=3, stride=1, padding=1, bias=False)
        self.rpe_bn = nn.BatchNorm2d(embed_dims)
        self.rpe_lif = MixedLIF(tau=2.0, detach_reset=True, backend='cupy')

    def forward(self, x):
        T, B, C, H, W = x.shape
        x = self.proj_conv(x.flatten(0, 1)) # have some fire value
        x = self.proj_bn(x).reshape(T, B, -1, H, W).contiguous()
        x = self.proj_lif(x).flatten(0, 1).contiguous()

        x = self.proj_conv1(x)
        x = self.proj_bn1(x).reshape(T, B, -1, H, W).contiguous()
        x = self.proj_lif1(x).flatten(0, 1).contiguous()

        x = self.proj_conv2(x)
        x = self.proj_bn2(x).reshape(T, B, -1, H, W).contiguous()
        x = self.proj_lif2(x).flatten(0, 1).contiguous()
        x = self.maxpool2(x)

        x = self.proj_conv3(x)
        x = self.proj_bn3(x).reshape(T, B, -1, H//2, W//2).contiguous()
        x = self.proj_lif3(x).flatten(0, 1).contiguous()
        x = self.maxpool3(x)

        x_feat = x.reshape(T, B, -1, H//4, W//4).contiguous()
        x = self.rpe_conv(x)
        x = self.rpe_bn(x).reshape(T, B, -1, H//4, W//4).contiguous()
        x = self.rpe_lif(x)
        x = x + x_feat

        x = x.flatten(-2).transpose(-1, -2)  # T,B,N,C， C=arg.dim， N=(32/4)^2
        return x


class Spikformer(nn.Module):
    def __init__(self,
                 img_size_h=128, img_size_w=128, patch_size=16, in_channels=2, num_classes=11,
                 embed_dims=[64, 128, 256], num_heads=[1, 2, 4], mlp_ratios=[4, 4, 4], qkv_bias=False, qk_scale=None,
                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm,
                 depths=[6, 8, 6], sr_ratios=[8, 4, 2], T = 4
                 ):
        super().__init__()
        self.T = T  # time step
        self.num_classes = num_classes
        self.depths = depths

        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depths)]  # stochastic depth decay rule

        patch_embed = SPS(img_size_h=img_size_h,
                                 img_size_w=img_size_w,
                                 patch_size=patch_size,
                                 in_channels=in_channels,
                                 embed_dims=embed_dims)

        block = nn.ModuleList([Block(
            dim=embed_dims, num_heads=num_heads, mlp_ratio=mlp_ratios, qkv_bias=qkv_bias,
            qk_scale=qk_scale, drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[j],
            norm_layer=norm_layer, sr_ratio=sr_ratios)
            for j in range(depths)])

        setattr(self, f"patch_embed", patch_embed)
        setattr(self, f"block", block)

        # classification head
        self.head = nn.Linear(embed_dims, num_classes) if num_classes > 0 else nn.Identity()
        self.apply(self._init_weights)

    @torch.jit.ignore
    def _get_pos_embed(self, pos_embed, patch_embed, H, W):
        if H * W == self.patch_embed1.num_patches:
            return pos_embed
        else:
            return F.interpolate(
                pos_embed.reshape(1, patch_embed.H, patch_embed.W, -1).permute(0, 3, 1, 2),
                size=(H, W), mode="bilinear").reshape(1, -1, H * W).permute(0, 2, 1)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def forward_features(self, x):

        block = getattr(self, f"block")
        patch_embed = getattr(self, f"patch_embed")

        x = patch_embed(x)
        for blk in block:
            x = blk(x)  # [T,B,N,C]
        return x.mean(2)  # [T, B, C]

    def forward(self, x):
        b_size = int(x.shape[0] / 2)
        x_temp = torch.zeros((self.T, b_size * 2,) + x.shape[1:], device=x.device)
        for t in range(self.T):
            x_temp[t, :b_size, ...] = x[:b_size, ...]
            x_temp[t, b_size:, ...] = x[b_size:, ...]

        x_f = self.forward_features(x_temp)  # [T, B, C]
        x = self.head(x_f.mean(0))

        x_f_view = torch.zeros((self.T * b_size * 2,) + x_f.shape[2:], device=x_f.device)
        for t in range(self.T):
            x_f_view[t * b_size:(t + 1) * b_size, ...] = x_f[t, :b_size, ...]
            x_f_view[(t + self.T) * b_size:(t + 1 + self.T) * b_size, ...] = x_f[t, b_size:, ...]
        return x, x_f_view


@register_model
def spikformer(pretrained=False, **kwargs):
    model = Spikformer(
        # img_size_h=224, img_size_w=224,
        # patch_size=16, embed_dims=768, num_heads=12, mlp_ratios=4,
        # in_channels=3, num_classes=1000, qkv_bias=False,
        # norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=12, sr_ratios=1,
        **kwargs
    )
    model.default_cfg = _cfg()
    return model


# class MixedLIF_(nn.Module):
#     """
#     Activative function is different for two trail of contrastive learning.
#     trail-1: original LIf
#     trail-2: Relu-like continuous func.
#     """
#     def __init__(self, tau: float = 2.0, v_threshold: float = 1.0, v_reset: float = 0.0,  detach_reset: bool = False,
#                  backend='torch'):
#         super(MixedLIF_, self).__init__()
#         self.tau = tau
#         self.v_threshold = v_threshold
#
#     def forward(self, x):
#         T = x.shape[0]
#         B = x.shape[1]
#         bs = B//2
#         u = torch.zeros((bs,)+x.shape[2:], device=x.device)
#         u2 = torch.zeros((bs,)+x.shape[2:], device=x.device)
#         o = torch.zeros(x.shape, device=x.device)
#         for t in range(T):
#             u = (1/self.tau) * u * (1 - spikefunc(u, self.v_threshold)) + x[t, :bs, ...]
#             u2 = (1/self.tau) * u2 * (1 - spikefunc(u, self.v_threshold)) + x[t, bs:, ...]
#             o[t, :bs, ...] = spikefunc(u, self.v_threshold)  # Equivalent to union of all spikes
#             o[t, bs:, ...] = torch.clamp(u2, min=self.v_threshold-0.5, max=self.v_threshold+0.5)
#         return o


# class SpikeFunction(torch.autograd.Function):
#     @staticmethod
#     def forward(ctx, input, v_threshold):
#         ctx.save_for_backward(input)
#         ctx.v_threshold = v_threshold
#
#         output = torch.gt(input, v_threshold)
#         return output.float()
#
#     @staticmethod
#     def backward(ctx, grad_output):
#         input, = ctx.saved_tensors
#         v_threshold = ctx.v_threshold
#         grad_input = grad_output.clone()
#         hu = (abs(input - v_threshold) < (a/2)) / a
#         return grad_input * hu, None
#
# spikefunc = SpikeFunction.apply


class MixedLIF(nn.Module):
    """
    Activative function of two Paths in the way of contrastive learing.
    Path-1: original LIf
    Path-2: Relu-like continuous func.
    Vth and tau are set same as class spikingjelly.activation_based.neuron.LIFNode
    """
    def __init__(self, tau: float = 2.0, v_threshold: float = 1.0, v_reset: float = 0.0,  detach_reset: bool = False,
                 backend='torch'):
        super(MixedLIF, self).__init__()
        self.lif = MultiStepLIFNode(tau=2.0, v_threshold=v_threshold, detach_reset=detach_reset, backend=backend)
        self.ClampLu = MultiStepLIFNode(tau=2.0, v_threshold=v_threshold, detach_reset=detach_reset, backend=backend)

    def forward(self, x):
        B = x.shape[1]
        bs = B//2

        o = torch.zeros(x.shape, device=x.device)
        o[:, :bs, ...] = self.lif(x[:, :bs, ...])
        o[:, bs:, ...] = self.ClampLu(x[:, bs:, ...])
        return o


class CustomLIFNode(neuron.BaseNode):
    def __init__(self, v_threshold: float = 1.0, v_reset: float = 0.0, tau: float = 2.0, detach_reset: bool = False,
                 backend='torch'):
        super().__init__(v_threshold=v_threshold, v_reset=v_reset, surrogate_function=None, detach_reset=detach_reset)
        self.tau = tau
        self.v_threshold = v_threshold

    def neuronal_charge(self, x: torch.Tensor):
        self.v = (1 - 1 / self.tau) * self.v + 1 / self.tau * x

    def neuronal_fire(self):
        return torch.clamp(self.v, min=self.v_threshold - 0.5, max=self.v_threshold + 0.5)-(self.v_threshold - 0.5)

    def neuronal_reset(self, spike: torch.Tensor):
        """
        * :ref:`API in English <BaseNode.neuronal_reset-en>`

        .. _BaseNode.neuronal_reset-cn:

        根据当前神经元释放的脉冲，对膜电位进行重置。

        * :ref:`中文API <BaseNode.neuronal_reset-cn>`

        .. _BaseNode.neuronal_reset-en:


        Reset the membrane potential according to neurons' output spikes.
        """
        @staticmethod
        @torch.jit.script
        def jit_hard_reset(v: torch.Tensor, spike: torch.Tensor, v_reset: float):
            v = (1. - spike) * v + spike * v_reset
            return v

        @staticmethod
        @torch.jit.script
        def jit_soft_reset(v: torch.Tensor, spike: torch.Tensor, v_threshold: float):
            v = v - spike * v_threshold
            return v

        if self.detach_reset:
            spike_d = spike.detach()
        else:
            spike_d = spike
        if self.v_reset is None:
            # soft reset
            self.v = jit_soft_reset(self.v, spike_d, self.v_threshold)
        else:
            # hard reset
            self.v = jit_hard_reset(self.v, spike_d, self.v_reset)
