"""Pluto: a small dense decoder-only LLM (GQA + QK-norm + gated attention + logit soft-capping).

Self-contained: needs only torch and transformers. A port of the dense path of the original
training code; weights are numerically equivalent to it (see test_export_hf.py in the source repo).

    model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
"""

from typing import Optional, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import DynamicCache, GenerationMixin, PreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast

from .configuration_pluto import PlutoConfig


class PlutoRMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return F.rms_norm(x, (x.size(-1),), self.weight, self.eps)


def rope_cos_sin(position_ids: torch.Tensor, head_dim: int, theta: float):
    """(B, T) integer positions -> cos and sin of shape (B, 1, T, head_dim / 2), float32"""
    inv_freq = 1.0 / theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=position_ids.device) / head_dim)
    freqs = position_ids.float()[..., None] * inv_freq  # (B, T, D/2)
    return freqs.cos()[:, None], freqs.sin()[:, None]


def apply_rope(x, cos, sin):
    """Rotate (B, H, T, D) by position, pairing dimension i with i + D/2"""
    x1, x2 = x.float().chunk(2, dim=-1)
    out = torch.cat((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1)
    return out.type_as(x)


class PlutoAttention(nn.Module):
    def __init__(self, config: PlutoConfig, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        self.n_heads = config.num_attention_heads
        self.n_kv = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.rope_theta = config.rope_theta
        d = config.hidden_size
        self.q_proj = nn.Linear(d, self.n_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(d, self.n_kv * self.head_dim, bias=False)
        self.v_proj = nn.Linear(d, self.n_kv * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.n_heads * self.head_dim, d, bias=False)
        self.g_proj = nn.Linear(d, self.n_heads * self.head_dim, bias=False) if config.attn_gate else None
        self.q_norm = PlutoRMSNorm(self.head_dim, config.rms_norm_eps) if config.qk_norm else nn.Identity()
        self.k_norm = PlutoRMSNorm(self.head_dim, config.rms_norm_eps) if config.qk_norm else nn.Identity()

    def forward(self, x, position_ids, attn_mask, is_causal, past_key_values=None):
        b, t, _ = x.shape
        q = self.q_norm(self.q_proj(x).view(b, t, self.n_heads, self.head_dim)).transpose(1, 2)
        k = self.k_norm(self.k_proj(x).view(b, t, self.n_kv, self.head_dim)).transpose(1, 2)
        v = self.v_proj(x).view(b, t, self.n_kv, self.head_dim).transpose(1, 2)
        cos, sin = rope_cos_sin(position_ids, self.head_dim, self.rope_theta)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        if past_key_values is not None:
            k, v = past_key_values.update(k, v, self.layer_idx)  # cached un-repeated, post-RoPE
        if self.n_kv != self.n_heads:
            rep = self.n_heads // self.n_kv
            k, v = k.repeat_interleave(rep, dim=1), v.repeat_interleave(rep, dim=1)
        y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, is_causal=is_causal)
        y = y.transpose(1, 2).reshape(b, t, -1)
        if self.g_proj is not None:
            y = y * torch.sigmoid(self.g_proj(x))
        return self.o_proj(y)


class PlutoMLP(nn.Module):
    def __init__(self, config: PlutoConfig):
        super().__init__()
        self.clamp = config.swiglu_clamp
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, x):
        gate, up = self.gate_proj(x), self.up_proj(x)
        if self.clamp is not None:
            gate, up = gate.clamp(max=self.clamp), up.clamp(-self.clamp, self.clamp)
        return self.down_proj(F.silu(gate) * up)


class PlutoDecoderLayer(nn.Module):
    def __init__(self, config: PlutoConfig, layer_idx: int):
        super().__init__()
        self.input_layernorm = PlutoRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.self_attn = PlutoAttention(config, layer_idx)
        self.post_attention_layernorm = PlutoRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mlp = PlutoMLP(config)

    def forward(self, x, position_ids, attn_mask, is_causal, past_key_values=None):
        x = x + self.self_attn(self.input_layernorm(x), position_ids, attn_mask, is_causal, past_key_values)
        return x + self.mlp(self.post_attention_layernorm(x))


class PlutoPreTrainedModel(PreTrainedModel):
    config_class = PlutoConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False
    _no_split_modules = ["PlutoDecoderLayer"]
    _supports_sdpa = True

    @torch.no_grad()
    def _init_weights(self, module):
        std = 0.02
        if isinstance(module, nn.Linear):
            module.weight.normal_(mean=0.0, std=std)
        elif isinstance(module, nn.Embedding):
            module.weight.normal_(mean=0.0, std=std)
        elif isinstance(module, PlutoRMSNorm):
            module.weight.fill_(1.0)


class PlutoModel(PlutoPreTrainedModel):
    def __init__(self, config: PlutoConfig):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(PlutoDecoderLayer(config, i) for i in range(config.num_hidden_layers))
        self.norm = PlutoRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[DynamicCache] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        use_cache: Optional[bool] = None,
        **kwargs,
    ) -> BaseModelOutputWithPast:
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Specify exactly one of input_ids and inputs_embeds")
        use_cache = self.config.use_cache if use_cache is None else use_cache
        x = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds
        b, t, _ = x.shape

        if use_cache and past_key_values is None:
            try:
                past_key_values = DynamicCache(config=self.config)
            except TypeError:  # older transformers
                past_key_values = DynamicCache()
        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
        cache_position = torch.arange(past_len, past_len + t, device=x.device)
        if position_ids is None:
            position_ids = cache_position[None].expand(b, -1)

        # Boolean mask, True = may attend; None lets SDPA take its causal fast paths.
        total = past_len + t
        padded = attention_mask is not None and attention_mask.shape[-1] == total and not bool(attention_mask.all())
        if not padded and t == 1:
            attn_mask, is_causal = None, False
        elif not padded and past_len == 0:
            attn_mask, is_causal = None, True
        else:
            keys = torch.arange(total, device=x.device)
            attn_mask = (keys[None, :] <= cache_position[:, None])[None, None]  # (1, 1, T, total)
            if padded:
                attn_mask = attn_mask & attention_mask.bool()[:, None, None, :]
                # Query rows that are all padding would softmax over nothing (NaN); let them see everything
                attn_mask = attn_mask | ~attn_mask.any(dim=-1, keepdim=True)
            is_causal = False

        for layer in self.layers:
            x = layer(x, position_ids, attn_mask, is_causal, past_key_values)
        return BaseModelOutputWithPast(
            last_hidden_state=self.norm(x), past_key_values=past_key_values if use_cache else None
        )


class PlutoForCausalLM(PlutoPreTrainedModel, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}

    def __init__(self, config: PlutoConfig):
        super().__init__(config)
        self.model = PlutoModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[DynamicCache] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        out = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
        )
        h = out.last_hidden_state
        if isinstance(logits_to_keep, int):
            keep = slice(-logits_to_keep, None) if logits_to_keep else slice(None)
        else:
            keep = logits_to_keep  # tensor of positions
        logits = self.lm_head(h[:, keep, :]).float()
        cap = self.config.logit_softcap
        if cap:
            logits = cap * torch.tanh(logits / cap)

        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1].reshape(-1, logits.size(-1))
            loss = F.cross_entropy(shift_logits, labels[:, 1:].reshape(-1).to(logits.device), ignore_index=-100)
        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=out.past_key_values)