from copy import deepcopy
from functools import partialmethod

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.jit import Final

from timm.layers import Mlp, DropPath
from timm.models.vision_transformer import LayerScale


class Linear(nn.Linear):

    def __init__(self, n_features, out_features, bias=True, device=None, dtype=None):

        super(Linear, self).__init__(n_features, out_features, bias=bias, device=device, dtype=dtype)
    
    def forward(self, *inputs):
        module_inputs, forward_modifiers = inputs

        assert isinstance(module_inputs, tuple), "Module inputs must be a tuple"
    
        execute_op = forward_modifiers.get("execute_op", None)
        if execute_op is None:
            raise ValueError("execute_op not found in forward_modifiers")
        
        primitive = forward_modifiers.get("primitive", None)
        if primitive is None:
            raise ValueError("primitive not found in forward_modifiers")
        
        if execute_op:
            # Reuse, Adapt
            x = super().forward(*module_inputs)
        else:
            # Skip, New
            x = module_inputs[0] if len(module_inputs) == 1 else module_inputs
        
        return x, primitive

    def __repr__(self):
        torch_repr = super().__repr__()
        return f"Conditional{torch_repr}"


class SeparableQKV(nn.Module):

    def __init__(self, dim, num_heads, head_dim, bias=True, device=None, dtype=None, artihippo_component=None):

        super().__init__()

        self.artihippo_component = artihippo_component or "value"

        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = head_dim
        _class = Linear if self.artihippo_component == "query" else nn.Linear
        self.q = _class(dim, dim, bias=bias, device=device, dtype=dtype)
        _class = Linear if self.artihippo_component == "key" else nn.Linear
        self.k = _class(dim, dim, bias=bias, device=device, dtype=dtype)
        _class = Linear if self.artihippo_component == "value" else nn.Linear
        self.v = _class(dim, dim, bias=bias, device=device, dtype=dtype)

        self.alias = {"query": "q", "key": "k", "value": "v"}

        self._weights_copied = False

    def copy_weights(self, weights, bias):

        if not self._weights_copied:
            with torch.no_grad():
                ptr = 0
                for component in ["q", "k", "v"]:
                    _weight = weights[ptr:ptr+self.dim, :]
                    _bias = bias[ptr:ptr+self.dim] if bias is not None else None
                    getattr(self, component).weight.copy_(_weight)
                    if bias is not None:
                        getattr(self, component).bias.copy_(_bias)
                    ptr += self.dim

            self._weights_copied = True
    
    def forward(self, x):
        
        assert self._weights_copied

        B, N, C = x.shape

        # The component to execute first is determined by the artihippo_component
        artihippo_component = self.alias[self.artihippo_component]
        other_components = [component for component in self.alias.values() if component != artihippo_component]

        artihippo_out, primitive = getattr(self, artihippo_component)(x)
        if primitive == "skip":
            return None, None, artihippo_out, primitive

        qkv = {artihippo_component: artihippo_out.reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)}

        for component in other_components:
            _component = getattr(self, component)(x)
            qkv[component] = _component.reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)

        return qkv["q"], qkv["k"], qkv["v"], primitive
    
    def __repr__(self):
        torch_repr = super().__repr__()
        return f"{self.artihippo_component}-{torch_repr}"


class ArtiHippoQKVAttention(nn.Module):
    fast_attn: Final[bool]

    def __init__(
            self,
            dim,
            num_heads=8,
            qkv_bias=False,
            qk_norm=False,
            attn_drop=0.,
            proj_drop=0.,
            norm_layer=nn.LayerNorm,
            artihippo_component=None
    ):
        super().__init__()
        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim ** -0.5
        self.fast_attn = hasattr(torch.nn.functional, 'scaled_dot_product_attention')  # FIXME

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.separable_qkv = SeparableQKV(dim, self.num_heads, self.head_dim, bias=qkv_bias, artihippo_component=artihippo_component)
        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        B, N, C = x.shape
        self.separable_qkv.copy_weights(self.qkv.weight, self.qkv.bias)
        
        q, k, v, primitive = self.separable_qkv(x)

        if primitive == "skip":
            # v can actually be skipped q or k. For skip, the outputs will be zeros
            # and the components will not be utilized anyway.
            return v

        q, k = self.q_norm(q), self.k_norm(k)

        if self.fast_attn:
            x = F.scaled_dot_product_attention(
                q, k, v,
                dropout_p=self.attn_drop.p,
            )
        else:
            q = q * self.scale
            attn = q @ k.transpose(-2, -1)
            attn = attn.softmax(dim=-1)
            attn = self.attn_drop(attn)
            x = attn @ v

        x = x.transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x
    

class ArtiHippoQKVBlock(nn.Module):

    def __init__(
            self,
            dim,
            num_heads,
            mlp_ratio=4.,
            qkv_bias=False,
            qk_norm=False,
            drop=0.,
            attn_drop=0.,
            init_values=None,
            drop_path=0.,
            act_layer=nn.GELU,
            norm_layer=nn.LayerNorm,
            artihippo_component=None
    ):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = ArtiHippoQKVAttention(
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_norm=qk_norm,
            attn_drop=attn_drop,
            proj_drop=drop,
            norm_layer=norm_layer,
            artihippo_component=artihippo_component
        )
        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()

        self.norm2 = norm_layer(dim)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=int(dim * mlp_ratio),
            act_layer=act_layer,
            drop=drop,
        )
        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
        self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x):
        x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x))))
        x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
        return x


def _ArtiHippoQKVBlock(cls, artihippo_component):
    class NewCls(cls):
        __init__ = partialmethod(cls.__init__, artihippo_component=artihippo_component)
    return NewCls


ArtiHippoQueryBlock = _ArtiHippoQKVBlock(ArtiHippoQKVBlock, artihippo_component="query")
ArtiHippoKeyBlock = _ArtiHippoQKVBlock(ArtiHippoQKVBlock, artihippo_component="key")
ArtiHippoValueBlock = _ArtiHippoQKVBlock(ArtiHippoQKVBlock, artihippo_component="value")
