from abc import ABC, abstractmethod
from typing import Dict, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from executorch.examples.models.llama.attention import (
    Attention,
    AttentionMHA,
    ForwardOptions,
    register_attention,
)
from executorch.examples.models.llama.model_args import ModelArgs
from executorch.examples.models.llama.rope import Rope


_CacheMap = Dict[str, torch.Tensor]
# Key and value caches are kept separate so the key caches can be kept transposed.
_InputCacheState = Tuple[_CacheMap, _CacheMap]
_OutputCacheState = Tuple[_CacheMap, _CacheMap]


class StaticKVCache(nn.Module, ABC):
    def __init__(self, layer_id: int, head_id: int):
        super().__init__()
        self.layer_id = layer_id
        self.head_id = head_id

    @abstractmethod
    def update(
        self,
        new_data: torch.Tensor,
        in_cache_state: Optional[_InputCacheState],
        out_cache_state: Optional[_OutputCacheState],
    ) -> Tuple[torch.Tensor, Optional[_OutputCacheState]]:
        """
        Given input cache state and new keys/values, returns the combined keys/values
        and the updated the output cache state.
        """
        pass

    def cache_key(self) -> str:
        return self.calculate_cache_key(self.layer_id, self.head_id)

    @staticmethod
    def calculate_cache_key(layer_id: int, head_id: int) -> str:
        return f"l{layer_id},h{head_id}"

    @staticmethod
    def apply_update(cache, update, pos, style, transpose=False):
        """
        After inference, update the cache state for next iteration. The runtime needs to
        implement the same operation.
        """
        if style == "shift_pointer":
            if transpose:
                update_len = update.size(-1)
                updated = torch.roll(cache, -update_len, -1)
                updated[:, :, -update_len:] = update
            else:
                update_len = update.size(-2)
                updated = torch.roll(cache, -update_len, -2)
                updated[:, -update_len:, :] = update

        if style == "smart_mask":
            updated = torch.clone(cache)
            if transpose:
                update_len = update.size(-1)
                updated[:, :, pos : pos + update_len] = update
            else:
                update_len = update.size(-2)
                updated[:, pos : pos + update_len, :] = update

        return updated


class StaticKCache(StaticKVCache):
    def __init__(self, layer_id: int, head_id: int, transpose=False):
        """
        If transpose is True, key cache is kept in (batch, dim, seq_len), otherwise in
        (batch, seq_len, dim).
        """
        super().__init__(layer_id, head_id)
        self.transpose = transpose

    def update(
        self,
        new_data: torch.Tensor,
        in_cache_state: Optional[_InputCacheState],
        out_cache_state: Optional[_OutputCacheState],
    ) -> Tuple[torch.Tensor, Optional[_OutputCacheState]]:
        seq_dim = -2
        if self.transpose:
            seq_dim = -1
            new_data = new_data.transpose(-1, -2)
        if in_cache_state is None:
            return new_data, None
        if out_cache_state is None:
            out_cache_state = ({}, {})

        all_data = torch.cat(
            [in_cache_state[0][self.cache_key()], new_data], dim=seq_dim
        )
        out_k_cache, out_v_cache = out_cache_state
        out_k_cache[self.cache_key()] = new_data
        return all_data, (out_k_cache, out_v_cache)


class StaticVCache(StaticKVCache):
    def update(
        self,
        new_data: torch.Tensor,
        in_cache_state: Optional[_InputCacheState],
        out_cache_state: Optional[_OutputCacheState],
    ) -> Tuple[torch.Tensor, Optional[_OutputCacheState]]:
        if in_cache_state is None:
            return new_data, None
        if out_cache_state is None:
            out_cache_state = ({}, {})

        all_data = torch.cat([in_cache_state[1][self.cache_key()], new_data], dim=-2)
        out_k_cache, out_v_cache = out_cache_state
        out_v_cache[self.cache_key()] = new_data
        return all_data, (out_k_cache, out_v_cache)


class StaticAttentionMask:
    def __init__(self, input_len, cache_len, style, mask_val=float("-inf")):
        self.input_len = input_len
        self.cache_len = cache_len
        assert style in ("shift_pointer", "smart_mask")
        self.style = style
        self.mask_val = mask_val
        self.unmasked_len = 0
        self.tensor = torch.zeros(1, input_len, input_len + cache_len)
        self.reset()

    def reset(self):
        self.unmasked_len = 0
        self.tensor[:, :, : self.cache_len] = self.mask_val

    def unmask(self, new_unmasked_len):
        if new_unmasked_len <= 0:
            return

        if self.style == "shift_pointer":
            self.tensor[
                :,
                :,
                self.cache_len
                - self.unmasked_len
                - new_unmasked_len : self.cache_len
                - self.unmasked_len,
            ] = 0

        if self.style == "smart_mask":
            self.tensor[
                :,
                :,
                self.unmasked_len : self.unmasked_len + new_unmasked_len,
            ] = 0

        self.unmasked_len += new_unmasked_len


class _Rope(nn.Module):
    def __init__(self, use_hf_rope):
        super().__init__()
        self.use_hf_rope = use_hf_rope

    def forward(
        self, x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor
    ) -> torch.Tensor:
        if self.use_hf_rope:
            if len(freqs_cos.shape) == 2:
                freqs_cos = freqs_cos.unsqueeze(0)
            if len(freqs_sin.shape) == 2:
                freqs_sin = freqs_sin.unsqueeze(0)
            x1 = x[..., : x.shape[-1] // 2]
            x2 = x[..., x.shape[-1] // 2 :]
            x_rotated = torch.cat((-x2, x1), dim=-1)
            return x * freqs_cos + x_rotated * freqs_sin
        else:
            x_r, x_i = x[..., ::2], x[..., 1::2]
            x_out_r = x_r * freqs_cos - x_i * freqs_sin
            x_out_i = x_r * freqs_sin + x_i * freqs_cos
            x_out = torch.cat([x_out_r, x_out_i], dim=-1)
            return x_out


@register_attention("static")
class StaticAttention(Attention):
    """
    An attention implementation meant for NPUs that require static shapes and are not
    flexible with tensor operations needed to perform KV cache updates. MHA/GQA is
    implemented as multiple SHAs, and the KV caches keep valid data at the end so the
    model only needs to perform a concat to combine past and new data.
    """

    def __init__(self, config: ModelArgs, layer_id: int, rope: Rope):
        super().__init__()
        self.n_heads = config.n_heads
        self.n_kv_heads = (
            self.n_heads if config.n_kv_heads is None else config.n_kv_heads
        )
        assert self.n_heads % self.n_kv_heads == 0
        self.n_heads_per_kv_group = self.n_heads // self.n_kv_heads
        self.dim = config.dim
        self.head_dim = config.head_dim
        self.inv_scale = 1.0 / (float(self.head_dim) ** 0.5)
        self.attention_qkv_bias = config.attention_qkv_bias
        self.use_qk_norm = config.use_qk_norm
        self.use_conv2d = False

        self.wqs = nn.ModuleList(
            [
                nn.Linear(self.dim, self.head_dim, bias=self.attention_qkv_bias)
                for _ in range(self.n_heads)
            ]
        )
        self.wks = nn.ModuleList(
            [
                nn.Linear(self.dim, self.head_dim, bias=self.attention_qkv_bias)
                for _ in range(self.n_kv_heads)
            ]
        )
        self.wvs = nn.ModuleList(
            [
                nn.Linear(self.dim, self.head_dim, bias=self.attention_qkv_bias)
                for _ in range(self.n_kv_heads)
            ]
        )

        self.k_caches = nn.ModuleList(
            [StaticKCache(layer_id, i) for i in range(self.n_kv_heads)]
        )
        self.v_caches = nn.ModuleList(
            [StaticVCache(layer_id, i) for i in range(self.n_kv_heads)]
        )
        self.wo = nn.Linear(self.n_heads * self.head_dim, self.dim, bias=False)
        self.rope = _Rope(rope.params.use_hf_rope)

        if self.use_qk_norm:
            self.q_norm = torch.nn.RMSNorm(self.head_dim, config.norm_eps)
            self.k_norm = torch.nn.RMSNorm(self.head_dim, config.norm_eps)
        else:
            self.q_norm = torch.nn.Identity()
            self.k_norm = torch.nn.Identity()

    def forward(
        self,
        x: torch.Tensor,
        freqs_cos: torch.Tensor,
        freqs_sin: torch.Tensor,
        **kwargs: ForwardOptions,
    ):
        mask = kwargs.get("mask")
        if (freqs_cos_override := kwargs.get("freqs_cos_override")) is not None:
            freqs_cos = freqs_cos_override  # pyre-ignore
        if (freqs_sin_override := kwargs.get("freqs_sin_override")) is not None:
            freqs_sin = freqs_sin_override  # pyre-ignore
        in_cache_state = kwargs.get("in_cache_state")
        out_cache_state = kwargs.get("out_cache_state")

        bsz, seq_len, dim = x.shape
        if self.use_conv2d:
            x = x.reshape(bsz, seq_len, 1, dim).transpose(1, 3)

        new_qs = [self.wqs[i](x) for i in range(self.n_heads)]
        new_ks = [self.wks[i](x) for i in range(self.n_kv_heads)]
        new_vs = [self.wvs[i](x) for i in range(self.n_kv_heads)]

        if self.use_conv2d:

            def from_conv2ds(ts):
                return [
                    t.reshape(bsz, self.head_dim, seq_len).transpose(1, 2) for t in ts
                ]

            new_qs = from_conv2ds(new_qs)
            new_ks = from_conv2ds(new_ks)
            new_vs = from_conv2ds(new_vs)

        if self.use_qk_norm:
            new_qs = [self.q_norm(q) for q in new_qs]
            new_ks = [self.k_norm(k) for k in new_ks]

        new_qs = [self.rope(q, freqs_cos, freqs_sin) for q in new_qs]
        new_ks = [self.rope(k, freqs_cos, freqs_sin) for k in new_ks]
        all_ks = []
        all_vs = []
        for i in range(self.n_kv_heads):
            ks, out_cache_state = self.k_caches[i].update(
                new_ks[i], in_cache_state, out_cache_state
            )
            all_ks.append(ks)
            vs, out_cache_state = self.v_caches[i].update(
                new_vs[i], in_cache_state, out_cache_state
            )
            all_vs.append(vs)

        heads = []
        for i in range(self.n_heads):
            kv_idx = i // self.n_heads_per_kv_group
            attn = new_qs[i] @ all_ks[kv_idx].transpose(-2, -1)
            attn = attn * self.inv_scale
            attn = attn + mask
            attn = F.softmax(attn, dim=-1)
            heads.append(attn @ all_vs[kv_idx])

        y = torch.cat(heads, dim=-1)
        if self.use_conv2d:
            y = (
                self.wo(y.reshape(bsz, seq_len, 1, -1).transpose(1, 3))
                .transpose(1, 3)
                .reshape(bsz, seq_len, -1)
            )
        else:
            y = self.wo(y)
        return y, {"out_cache_state": out_cache_state}

    def load_weights_from_attention_mha(self, other: AttentionMHA):
        for i in range(self.n_heads):
            self.wqs[i].weight.data.copy_(
                other.wq.weight[i * self.head_dim : (i + 1) * self.head_dim, :]
            )

        for i in range(self.n_kv_heads):
            self.wks[i].weight.data.copy_(
                other.wk.weight[i * self.head_dim : (i + 1) * self.head_dim, :]
            )
            self.wvs[i].weight.data.copy_(
                other.wv.weight[i * self.head_dim : (i + 1) * self.head_dim, :]
            )

        self.wo.weight.data.copy_(other.wo.weight)

        if other.use_qk_norm:
            self.use_qk_norm = True
            self.q_norm = torch.nn.RMSNorm(other.q_norm_fn.dim, other.q_norm_fn.eps)
            self.q_norm.load_state_dict(other.q_norm_fn.state_dict())
            self.k_norm = torch.nn.RMSNorm(other.k_norm_fn.dim, other.k_norm_fn.eps)
            self.k_norm.load_state_dict(other.k_norm_fn.state_dict())

    def linear_to_conv2d(self):
        def transfer_weight(linear, conv2d):
            conv2d.weight.data.copy_(linear.weight[:, :, None, None])
            return conv2d

        self.wqs = nn.ModuleList(
            [
                transfer_weight(
                    linear,
                    nn.Conv2d(self.dim, self.head_dim, 1, bias=self.attention_qkv_bias),
                )
                for linear in self.wqs
            ]
        )
        self.wks = nn.ModuleList(
            [
                transfer_weight(
                    linear,
                    nn.Conv2d(self.dim, self.head_dim, 1, bias=self.attention_qkv_bias),
                )
                for linear in self.wks
            ]
        )
        self.wvs = nn.ModuleList(
            [
                transfer_weight(
                    linear,
                    nn.Conv2d(self.dim, self.head_dim, 1, bias=self.attention_qkv_bias),
                )
                for linear in self.wvs
            ]
        )
        self.wo = transfer_weight(
            self.wo,
            nn.Conv2d(
                self.n_heads * self.head_dim, self.dim, 1, bias=self.attention_qkv_bias
            ),
        )

        self.use_conv2d = True