"""Portable PyTorch Weevil, with fp32 recurrence and fixed-size generation caches.

This reference implementation needs only PyTorch and Transformers, including for
the latent-QKV variant. Both embedding factors are tied to the vocabulary readout.
"""

from typing import ClassVar

import torch
from torch import nn
from torch.nn import functional as F
from transformers import GenerationMixin, PreTrainedModel
from transformers.cache_utils import Cache
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_weevil import WeevilConfig


class WeevilCache(Cache):
    def __init__(self, states=None, seen_tokens=0):
        super().__init__(layers=[])
        self.states = states or []
        self.seen_tokens = seen_tokens

    def get_seq_length(self, layer_idx=0):
        return self.seen_tokens

    def reorder_cache(self, beam_idx):
        self.states = [tuple(x.index_select(0, beam_idx.to(x.device)) for x in state)
                       for state in self.states]

    def batch_repeat_interleave(self, repeats):
        self.states = [tuple(x.repeat_interleave(repeats, dim=0) for x in state)
                       for state in self.states]


class ZCRMSNorm(nn.Module):
    def __init__(self, width, eps):
        super().__init__()
        self.scale = nn.Parameter(torch.zeros(width))
        self.eps = eps

    def forward(self, x):
        x = x.float()
        return (1 + self.scale.float()) * x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps)


class TiedEmbedding(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.embed = nn.Embedding(c.vocab_size, c.embedding_dim)
        self.proj = nn.Linear(c.embedding_dim, c.hidden_size, bias=False)

    def forward(self, ids):
        return self.proj(self.embed(ids))

    def token_logits(self, x):
        x = F.linear(x.to(self.proj.weight.dtype), self.proj.weight.T)
        return F.linear(x, self.embed.weight).float()


class CausalConv(nn.Module):
    def __init__(self, channels, kernel):
        super().__init__()
        self.weight = nn.Parameter(torch.empty(channels, kernel))
        nn.init.normal_(self.weight, std=0.02)

    def step(self, x, history, valid):
        taps = torch.cat((history, x.float().unsqueeze(1)), dim=1)
        output = (taps * self.weight.float().T).sum(1)
        history = torch.where(valid[:, None, None], taps[:, 1:], history)
        return F.silu(output.to(x.dtype)), history


class GDN2Mixer(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.config = c
        hidden, head = c.hidden_size, c.hidden_size // c.num_heads
        width = c.num_heads * head

        def linear(i, o, bias=False):
            return nn.Linear(i, o, bias=bias)

        if c.qkv_rank is None:
            self.q_proj, self.k_proj, self.v_proj = [linear(hidden, width) for _ in range(3)]
        else:
            self.qkv_a = linear(hidden, c.qkv_rank)
            self.q_b, self.k_b, self.v_b = [linear(c.qkv_rank, width) for _ in range(3)]
        self.q_conv, self.k_conv, self.v_conv = [CausalConv(width, c.conv_kernel) for _ in range(3)]
        self.f_in, self.f_out = linear(hidden, head), linear(head, width)
        self.b_proj, self.w_proj = linear(hidden, width), linear(hidden, width)
        self.g_in, self.g_out = linear(hidden, head), linear(head, width, bias=True)
        self.o_proj = linear(width, hidden)
        self.o_norm = nn.Parameter(torch.ones(head))
        self.A_log = nn.Parameter(torch.zeros(c.num_heads))
        self.dt_bias = nn.Parameter(torch.zeros(width))

    def forward(self, x, valid, state=None):
        c = self.config
        batch, length, _ = x.shape
        head = c.hidden_size // c.num_heads
        width = c.num_heads * head
        if state is None:
            recurrent = torch.zeros(batch, c.num_heads, head, head, device=x.device)
            histories = [torch.zeros(batch, c.conv_kernel - 1, width, device=x.device) for _ in range(3)]
        else:
            recurrent, *histories = state
        if c.qkv_rank is None:
            projections = [self.q_proj(x), self.k_proj(x), self.v_proj(x)]
        else:
            latent = self.qkv_a(x)
            projections = [self.q_b(latent), self.k_b(latent), self.v_b(latent)]
        shape = (batch, length, c.num_heads, head)
        decay = -self.A_log.float().exp()[None, None, :, None] * F.softplus(
            self.f_out(self.f_in(x)).float() + self.dt_bias.float()).reshape(shape)
        erase = self.b_proj(x).sigmoid().float().reshape(shape)
        write = self.w_proj(x).sigmoid().float().reshape(shape)
        outputs = []
        for t in range(length):
            active = valid[:, t]
            qkv = []
            for i, conv in enumerate((self.q_conv, self.k_conv, self.v_conv)):
                value, histories[i] = conv.step(projections[i][:, t], histories[i], active)
                qkv.append(value.float().reshape(batch, c.num_heads, head))
            q, k, v = qkv
            q = q * torch.rsqrt(q.square().sum(-1, keepdim=True) + 1e-6) * head ** -0.5
            k = k * torch.rsqrt(k.square().sum(-1, keepdim=True) + 1e-6)
            decayed = recurrent * decay[:, t].exp().unsqueeze(-1)
            erased = ((erase[:, t] * k).unsqueeze(-1) * decayed).sum(-2)
            updated = decayed + k.unsqueeze(-1) * (write[:, t] * v - erased).unsqueeze(-2)
            recurrent = torch.where(active[:, None, None, None], updated, recurrent)
            outputs.append((q.unsqueeze(-1) * recurrent).sum(-2))
        output = torch.stack(outputs, dim=1)
        output = output * torch.rsqrt(output.square().mean(-1, keepdim=True) + c.norm_eps)
        gate = self.g_out(self.g_in(x)).float().reshape(shape).sigmoid()
        output = output * self.o_norm.float() * gate
        output = self.o_proj(output.flatten(-2).to(x.dtype))
        return output.masked_fill(~valid[..., None], 0), (recurrent, *histories)


class SwiGLU(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.gate_proj = nn.Linear(c.hidden_size, c.intermediate_size, bias=False)
        self.up_proj = nn.Linear(c.hidden_size, c.intermediate_size, bias=False)
        self.down_proj = nn.Linear(c.intermediate_size, c.hidden_size, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class WeevilLayer(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.norm = ZCRMSNorm(c.hidden_size, c.norm_eps)
        self.mixer = GDN2Mixer(c)
        self.mlp = SwiGLU(c)

    def forward(self, x, valid, state):
        z = self.norm(x).to(self.mixer.o_proj.weight.dtype)
        mixed, state = self.mixer(z, valid, state)
        x = x.float() + mixed.float() + self.mlp(z).float()
        return x.masked_fill(~valid[..., None], 0), state


class WeevilForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = WeevilConfig
    base_model_prefix = "weevil"
    main_input_name = "input_ids"
    _no_split_modules: ClassVar[list[str]] = ["WeevilLayer"]

    def __init__(self, config):
        super().__init__(config)
        self.embedding = TiedEmbedding(config)
        self.decoder = nn.ModuleDict({"layers": nn.ModuleList([
            WeevilLayer(config) for _ in range(config.num_hidden_layers)
        ])})
        self.final_norm = ZCRMSNorm(config.hidden_size, config.norm_eps)
        self.post_init()

    @classmethod
    def _supports_default_dynamic_cache(cls):
        return False

    def get_input_embeddings(self):
        return self.embedding.embed

    def set_input_embeddings(self, value):
        self.embedding.embed = value

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None,
                                      attention_mask=None, use_cache=True, **kwargs):
        if past_key_values is not None and past_key_values.get_seq_length():
            input_ids = input_ids[:, past_key_values.get_seq_length():]
        return {"input_ids": input_ids, "attention_mask": attention_mask,
                "past_key_values": past_key_values, "use_cache": use_cache}

    def forward(self, input_ids=None, attention_mask=None, past_key_values=None,
                inputs_embeds=None, labels=None, use_cache=None, return_dict=None,
                token_type_ids=None, **kwargs):
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Provide exactly one of input_ids or inputs_embeds")
        x = self.embedding(input_ids) if inputs_embeds is None else inputs_embeds
        batch, length = x.shape[:2]
        valid = (torch.ones(batch, length, dtype=torch.bool, device=x.device)
                 if attention_mask is None else attention_mask[:, -length:].to(device=x.device, dtype=torch.bool))
        x = x.float().masked_fill(~valid[..., None], 0)
        use_cache = self.config.use_cache if use_cache is None else use_cache
        if past_key_values is not None and not isinstance(past_key_values, WeevilCache):
            raise ValueError("Weevil requires a WeevilCache, not an attention KV cache")
        states = []
        for i, layer in enumerate(self.decoder["layers"]):
            state = past_key_values.states[i] if past_key_values is not None and past_key_values.states else None
            x, state = layer(x, valid, state)
            states.append(state)
        logits = self.embedding.token_logits(self.final_norm(x))
        logits[..., self.config.pad_token_id] = -1e9
        logits = logits.masked_fill(~valid[..., None], 0)
        loss = None
        if labels is not None:
            targets = labels[:, 1:].to(logits.device).clone()
            targets.masked_fill_(~valid[:, 1:], -100)
            loss = F.cross_entropy(logits[:, :-1].reshape(-1, self.config.vocab_size),
                                   targets.reshape(-1), ignore_index=-100)
        seen = 0 if past_key_values is None else past_key_values.get_seq_length()
        cache = WeevilCache(states, seen + length) if use_cache else None
        if return_dict is False:
            result = (logits, cache) if use_cache else (logits,)
            return (loss, *result) if loss is not None else result
        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=cache)


WeevilForCausalLM.register_for_auto_class("AutoModelForCausalLM")