"""Hugging Face PreTrainedModel wrapper for SimpleLM (auto-generated).

Module structure mirrors the upstream `DecoderOnlyLM` exactly so the
state_dict in `model.safetensors` loads with no key remapping.
"""

from __future__ import annotations

from typing import Optional

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

from .configuration_simple_lm import SimpleLMConfig


def _activation_module(name: str) -> nn.Module:
    if name == "gelu":
        return nn.GELU(approximate="tanh")
    if name == "relu":
        return nn.ReLU()
    if name == "silu":
        return nn.SiLU()
    raise ValueError(f"Unknown activation: {name}")


class _CausalSelfAttention(nn.Module):
    def __init__(self, cfg: SimpleLMConfig) -> None:
        super().__init__()
        self.attn = nn.MultiheadAttention(
            cfg.d_model,
            cfg.n_heads,
            dropout=cfg.dropout,
            bias=cfg.bias,
            batch_first=True,
        )
        self.dropout = nn.Dropout(cfg.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        seq_len = x.size(1)
        mask = torch.triu(
            torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool),
            diagonal=1,
        )
        out, _ = self.attn(x, x, x, attn_mask=mask, need_weights=False)
        return self.dropout(out)


class _TransformerBlock(nn.Module):
    def __init__(self, cfg: SimpleLMConfig) -> None:
        super().__init__()
        self.ln1 = nn.LayerNorm(cfg.d_model, bias=cfg.bias)
        self.attn = _CausalSelfAttention(cfg)
        self.ln2 = nn.LayerNorm(cfg.d_model, bias=cfg.bias)
        ffn = cfg.d_ff if cfg.d_ff is not None else 4 * cfg.d_model
        self.mlp = nn.Sequential(
            nn.Linear(cfg.d_model, ffn, bias=cfg.bias),
            _activation_module(cfg.activation),
            nn.Dropout(cfg.dropout),
            nn.Linear(ffn, cfg.d_model, bias=cfg.bias),
            nn.Dropout(cfg.dropout),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x


class SimpleLMForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = SimpleLMConfig
    base_model_prefix = "simple_lm"
    main_input_name = "input_ids"
    _tied_weights_keys = {"lm_head.weight": "tok_emb.weight"}

    def __init__(self, cfg: SimpleLMConfig) -> None:
        super().__init__(cfg)
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.pos_emb = nn.Embedding(cfg.context_length, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList(
            _TransformerBlock(cfg) for _ in range(cfg.n_layers)
        )
        self.ln_f = nn.LayerNorm(cfg.d_model, bias=cfg.bias)
        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        if cfg.tie_word_embeddings:
            self.lm_head.weight = self.tok_emb.weight
        self.post_init()

    def get_input_embeddings(self) -> nn.Module:
        return self.tok_emb

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.tok_emb = value

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

    def set_output_embeddings(self, value: nn.Module) -> None:
        self.lm_head = value

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        ctx = self.config.context_length
        if input_ids.size(1) > ctx:
            input_ids = input_ids[:, -ctx:]
        b, t = input_ids.shape
        pos = torch.arange(t, device=input_ids.device).unsqueeze(0).expand(b, t)
        x = self.drop(self.tok_emb(input_ids) + self.pos_emb(pos))
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)
        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
            )
        return CausalLMOutputWithPast(loss=loss, logits=logits)

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> dict:
        ctx = self.config.context_length
        if input_ids.size(1) > ctx:
            input_ids = input_ids[:, -ctx:]
        return {"input_ids": input_ids, "attention_mask": attention_mask}