"""Portable Transformers implementation of the Maccy architecture."""

from collections.abc import Sequence
from typing import Any, cast

import torch
from torch import Tensor, nn
from torch.nn import functional as F
from transformers import PreTrainedModel
from transformers.generation.utils import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_maccy import MaccyConfig

_NORM_EPSILON = 1e-6
_MINIMUM_RETENTION = 0.125


class RMSNorm(nn.Module):
    """Normalize vector magnitude without subtracting its mean."""

    def __init__(self, width: int) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(width))

    def forward(self, inputs: Tensor) -> Tensor:
        inverse_rms = torch.rsqrt(
            inputs.float().square().mean(dim=-1, keepdim=True) + _NORM_EPSILON
        )
        return inputs * inverse_rms.to(inputs.dtype) * self.weight.to(inputs.dtype)


class RotaryEmbedding(nn.Module):
    """Apply rotary position embeddings over the penultimate dimension."""

    def __init__(self, width: int, maximum_length: int) -> None:
        super().__init__()
        self.width = width
        self.maximum_length = maximum_length

    def forward(self, inputs: Tensor) -> Tensor:
        sequence_length = inputs.shape[-2]
        if sequence_length > self.maximum_length:
            raise ValueError("sequence length exceeds the rotary embedding limit")
        pair_indices = torch.arange(0, self.width, 2, dtype=torch.float32, device=inputs.device)
        inverse_frequencies = 1.0 / (10_000.0 ** (pair_indices / self.width))
        positions = torch.arange(sequence_length, dtype=torch.float32, device=inputs.device)
        angles = torch.outer(positions, inverse_frequencies)
        cosines = angles.cos().to(inputs.dtype)
        sines = angles.sin().to(inputs.dtype)
        even, odd = inputs[..., 0::2], inputs[..., 1::2]
        return torch.stack(
            (even * cosines - odd * sines, even * sines + odd * cosines), dim=-1
        ).flatten(start_dim=-2)


def recurrent_kda(
    queries: Tensor,
    keys: Tensor,
    values: Tensor,
    retention: Tensor,
    update_rate: Tensor,
) -> Tensor:
    """Evaluate the delta-rule memory recurrence in float32."""
    output_dtype = values.dtype
    queries = queries.float() * (keys.shape[-1] ** -0.5)
    keys = keys.float()
    values = values.float()
    retention = retention.float()
    update_rate = update_rate.float()

    batch_size, _, n_heads, key_dim = keys.shape
    state = keys.new_zeros(batch_size, n_heads, key_dim, values.shape[-1])
    outputs = []
    for token_index in range(keys.shape[1]):
        query = queries[:, token_index]
        key = keys[:, token_index]
        value = values[:, token_index]
        state = state * retention[:, token_index].unsqueeze(-1)
        prediction = torch.einsum("bhkv,bhk->bhv", state, key)
        error = value - prediction
        beta = update_rate[:, token_index, :, None, None]
        state = state + beta * key.unsqueeze(-1) * error.unsqueeze(-2)
        outputs.append(torch.einsum("bhkv,bhk->bhv", state, query))
    return torch.stack(outputs, dim=1).to(output_dtype)


class CausalDepthwiseConvolution(nn.Module):
    """Mix a four-token local history independently within each channel."""

    def __init__(self, d_model: int) -> None:
        super().__init__()
        self.width = 4
        self.convolution = nn.Conv1d(
            d_model, d_model, kernel_size=self.width, groups=d_model, bias=False, padding=3
        )

    def forward(self, inputs: Tensor) -> Tensor:
        sequence_length = inputs.shape[1]
        convolved = self.convolution(inputs.transpose(1, 2))[..., :sequence_length]
        return F.silu(convolved.transpose(1, 2))


class KimiDeltaAttention(nn.Module):
    """Kimi Delta Attention with a portable recurrent implementation."""

    def __init__(self, d_model: int, n_heads: int) -> None:
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.qkv_projection = nn.Linear(d_model, 3 * d_model, bias=False)
        self.query_convolution = CausalDepthwiseConvolution(d_model)
        self.key_convolution = CausalDepthwiseConvolution(d_model)
        self.value_convolution = CausalDepthwiseConvolution(d_model)
        self.control_down = nn.Linear(d_model, 2 * self.head_dim, bias=False)
        self.update_projection = nn.Linear(d_model, n_heads, bias=False)
        self.retention_up = nn.Linear(self.head_dim, d_model, bias=False)
        self.retention_bias = nn.Parameter(torch.zeros(d_model))
        self.log_decay_scale = nn.Parameter(torch.zeros(n_heads))
        self.output_gate_up = nn.Linear(self.head_dim, d_model, bias=True)
        self.output_norm_weight = nn.Parameter(torch.ones(self.head_dim))
        self.output = nn.Linear(d_model, d_model, bias=False)

    def _split_heads(self, inputs: Tensor) -> Tensor:
        return inputs.view(inputs.shape[0], inputs.shape[1], self.n_heads, self.head_dim)

    def forward(self, inputs: Tensor) -> Tensor:
        queries, keys, values = self.qkv_projection(inputs).chunk(3, dim=-1)
        queries = F.normalize(
            self._split_heads(self.query_convolution(queries)), dim=-1, eps=_NORM_EPSILON
        )
        keys = F.normalize(self._split_heads(self.key_convolution(keys)), dim=-1, eps=_NORM_EPSILON)
        values = self._split_heads(self.value_convolution(values))

        retention_latent, gate_latent = self.control_down(inputs).chunk(2, dim=-1)
        retention_logits = self.retention_up(retention_latent) + self.retention_bias
        retention_logits = self._split_heads(retention_logits).float()
        decay_scale = self.log_decay_scale.exp().view(1, 1, self.n_heads, 1)
        retention = (-decay_scale * F.softplus(retention_logits)).exp()
        retention = retention.clamp_min(_MINIMUM_RETENTION)
        update_rate = self.update_projection(inputs).float().sigmoid()
        output_gate = self._split_heads(self.output_gate_up(gate_latent)).float().sigmoid()

        mixed = recurrent_kda(queries, keys, values, retention, update_rate).float()
        inverse_rms = torch.rsqrt(mixed.square().mean(dim=-1, keepdim=True) + _NORM_EPSILON)
        mixed = mixed * inverse_rms * self.output_norm_weight.float()
        mixed = mixed * output_gate
        mixed = mixed.to(self.output.weight.dtype)
        return self.output(mixed.flatten(start_dim=-2))


class MultiHeadLatentAttention(nn.Module):
    """Causal attention through compressed query and key-value latents."""

    def __init__(self, config: MaccyConfig) -> None:
        super().__init__()
        mla = config.mla
        self.n_heads = config.n_heads
        self.query_rank = mla["query_rank"]
        self.kv_rank = mla["kv_rank"]
        self.content_head_dim = mla["content_head_dim"]
        self.rope_head_dim = mla["rope_head_dim"]
        self.value_head_dim = mla["value_head_dim"]
        query_head_dim = self.content_head_dim + self.rope_head_dim

        self.input_down = nn.Linear(
            config.d_model,
            self.query_rank + self.kv_rank + self.rope_head_dim,
            bias=config.bias,
        )
        self.query_norm = RMSNorm(self.query_rank)
        self.query_up = nn.Linear(self.query_rank, self.n_heads * query_head_dim, bias=config.bias)
        self.kv_norm = RMSNorm(self.kv_rank)
        self.kv_up = nn.Linear(
            self.kv_rank,
            self.n_heads * (self.content_head_dim + self.value_head_dim),
            bias=config.bias,
        )
        self.rotary_embedding = RotaryEmbedding(self.rope_head_dim, config.context_length)
        self.gate = (
            nn.Linear(config.d_model, self.n_heads * self.value_head_dim, bias=True)
            if mla["gated"]
            else None
        )
        self.output = nn.Linear(
            self.n_heads * self.value_head_dim, config.d_model, bias=config.bias
        )

    def forward(self, inputs: Tensor) -> Tensor:
        batch_size, sequence_length, _ = inputs.shape
        compressed_queries, compressed_kv, rotary_keys = self.input_down(inputs).split(
            (self.query_rank, self.kv_rank, self.rope_head_dim), dim=-1
        )
        expanded_queries = self.query_up(self.query_norm(compressed_queries)).view(
            batch_size,
            sequence_length,
            self.n_heads,
            self.content_head_dim + self.rope_head_dim,
        )
        expanded_kv = self.kv_up(self.kv_norm(compressed_kv)).view(
            batch_size,
            sequence_length,
            self.n_heads,
            self.content_head_dim + self.value_head_dim,
        )
        content_queries, rotary_queries = expanded_queries.split(
            (self.content_head_dim, self.rope_head_dim), dim=-1
        )
        content_keys, values = expanded_kv.split(
            (self.content_head_dim, self.value_head_dim), dim=-1
        )
        rotary_queries = self.rotary_embedding(rotary_queries.transpose(1, 2))
        rotary_keys = self.rotary_embedding(rotary_keys.unsqueeze(1)).expand(
            -1, self.n_heads, -1, -1
        )
        queries = torch.cat((content_queries.transpose(1, 2), rotary_queries), dim=-1)
        keys = torch.cat((content_keys.transpose(1, 2), rotary_keys), dim=-1)
        values = values.transpose(1, 2)
        mixed = F.scaled_dot_product_attention(queries, keys, values, is_causal=True)
        mixed = mixed.transpose(1, 2)
        if self.gate is not None:
            gate = self.gate(inputs).view(
                batch_size, sequence_length, self.n_heads, self.value_head_dim
            )
            mixed = mixed * gate.sigmoid()
        return self.output(mixed.flatten(start_dim=-2))


class PackedSwiGLUExperts(nn.Module):
    """Store equal-shaped experts in two packed parameter tensors."""

    def __init__(self, n_experts: int, d_model: int, hidden_dim: int, *, bias: bool) -> None:
        super().__init__()
        self.input_weight = nn.Parameter(torch.empty(n_experts, d_model, 2 * hidden_dim))
        self.output_weight = nn.Parameter(torch.empty(n_experts, hidden_dim, d_model))
        if bias:
            self.input_bias = nn.Parameter(torch.zeros(n_experts, 2 * hidden_dim))
            self.output_bias = nn.Parameter(torch.zeros(n_experts, d_model))
        else:
            self.register_parameter("input_bias", None)
            self.register_parameter("output_bias", None)

    def forward_expert(self, expert_index: int, inputs: Tensor) -> Tensor:
        projected = inputs @ self.input_weight[expert_index]
        if self.input_bias is not None:
            projected = projected + self.input_bias[expert_index]
        gate, values = projected.chunk(2, dim=-1)
        updates = (F.silu(gate) * values) @ self.output_weight[expert_index]
        if self.output_bias is not None:
            updates = updates + self.output_bias[expert_index]
        return updates


class SparseMoE(nn.Module):
    """Route each token to a weighted top-k subset of SwiGLU experts."""

    def __init__(self, config: MaccyConfig) -> None:
        super().__init__()
        moe = config.moe
        self.n_experts = moe["n_experts"]
        self.experts_per_token = moe["experts_per_token"]
        self.router = nn.Linear(config.d_model, self.n_experts, bias=False)
        hidden_dim = round(moe["expert_expansion"] * config.d_model)
        self.experts = PackedSwiGLUExperts(
            self.n_experts, config.d_model, hidden_dim, bias=config.bias
        )

    def forward(self, inputs: Tensor) -> Tensor:
        input_shape = inputs.shape
        flat_inputs = inputs.flatten(0, -2)
        probabilities = self.router(flat_inputs).float().softmax(dim=-1)
        weights, expert_indices = probabilities.topk(self.experts_per_token, dim=-1)
        weights = (weights / weights.sum(dim=-1, keepdim=True)).to(inputs.dtype)
        updates = torch.zeros_like(flat_inputs)
        for expert_index in range(self.n_experts):
            token_indices, choice_indices = torch.where(expert_indices == expert_index)
            if token_indices.numel() == 0:
                continue
            expert_updates = self.experts.forward_expert(
                expert_index, flat_inputs.index_select(0, token_indices)
            )
            expert_weights = weights[token_indices, choice_indices].unsqueeze(-1)
            updates = updates.index_add(0, token_indices, expert_updates * expert_weights)
        return updates.view(input_shape)


class TransformerBlock(nn.Module):
    """Apply one pre-normalized sequence mixer and sparse channel mixer."""

    def __init__(self, config: MaccyConfig, mixer_kind: str) -> None:
        super().__init__()
        self.attention_norm = RMSNorm(config.d_model)
        self.mixer = (
            KimiDeltaAttention(config.d_model, config.n_heads)
            if mixer_kind == "kda"
            else MultiHeadLatentAttention(config)
        )
        self.feed_forward_norm = RMSNorm(config.d_model)
        self.feed_forward = SparseMoE(config)

    def forward(self, inputs: Tensor) -> Tensor:
        inputs = inputs + self.mixer(self.attention_norm(inputs))
        return inputs + self.feed_forward(self.feed_forward_norm(inputs))


class MaccyPreTrainedModel(PreTrainedModel):
    """Shared Transformers metadata for Maccy models."""

    config_class = MaccyConfig
    base_model_prefix = ""
    _no_split_modules = ["TransformerBlock"]
    _supports_sdpa = True

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, (nn.Linear, nn.Embedding, nn.Conv1d)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                nn.init.zeros_(module.bias)


class MaccyForCausalLM(MaccyPreTrainedModel, GenerationMixin):
    """Maccy decoder with a tied next-token language-modeling head."""

    _tied_weights_keys = {"lm_head.weight": "token_embedding.weight"}

    def __init__(self, config: MaccyConfig) -> None:
        super().__init__(config)
        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
        mixers = tuple(part.strip() for part in config.mixer_pattern.split(","))
        repeated_mixers = mixers * (config.n_layers // len(mixers))
        self.blocks = nn.ModuleList(
            TransformerBlock(config, mixer_kind) for mixer_kind in repeated_mixers
        )
        self.output_norm = RMSNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=config.bias)
        self.post_init()

    def get_input_embeddings(self) -> nn.Embedding:
        return self.token_embedding

    def set_input_embeddings(self, value: nn.Module) -> None:
        if not isinstance(value, nn.Embedding):
            raise TypeError("input embeddings must be an nn.Embedding")
        self.token_embedding = value

    def get_output_embeddings(self) -> nn.Linear:
        return self.lm_head

    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
        if not isinstance(new_embeddings, nn.Linear):
            raise TypeError("output embeddings must be an nn.Linear")
        self.lm_head = new_embeddings

    def forward(
        self,
        input_ids: Tensor | None = None,
        attention_mask: Tensor | None = None,
        inputs_embeds: Tensor | None = None,
        labels: Tensor | None = None,
        use_cache: bool | None = None,
        logits_to_keep: int | Tensor = 0,
        return_dict: bool | None = None,
        **_: Any,
    ) -> CausalLMOutputWithPast | tuple[Tensor, ...]:
        del attention_mask, use_cache
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("pass exactly one of input_ids or inputs_embeds")
        hidden_states = self.token_embedding(input_ids) if inputs_embeds is None else inputs_embeds
        if hidden_states.shape[1] > self.config.context_length:
            raise ValueError(f"Maccy's context length is {self.config.context_length} tokens")
        for block in self.blocks:
            hidden_states = block(hidden_states)
        hidden_states = self.output_norm(hidden_states)
        indices = (
            slice(None)
            if labels is not None or (isinstance(logits_to_keep, int) and logits_to_keep == 0)
            else slice(-logits_to_keep, None)
            if isinstance(logits_to_keep, int)
            else logits_to_keep
        )
        logits = self.lm_head(hidden_states[:, indices, :])

        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1].contiguous().float()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, self.config.vocab_size),
                shift_labels.view(-1),
                ignore_index=-100,
            )

        output = CausalLMOutputWithPast(
            loss=cast(torch.FloatTensor | None, loss),
            logits=logits,
            past_key_values=None,
        )
        if return_dict is False:
            values: Sequence[Tensor | None] = (loss, logits) if loss is not None else (logits,)
            return tuple(value for value in values if value is not None)
        return output