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

from .configuration_tinygpt import TinyGPTConfig


# ─────────────────────────────────────────────
# These three classes are copied verbatim (module
# names included) from the original training script,
# so a raw_model.state_dict() from training loads
# straight into TinyGPTForCausalLM with strict=True.
# ─────────────────────────────────────────────
class MultiHeadAttention(nn.Module):
    def __init__(self, n_embd, n_head):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.head_size = n_embd // n_head
        self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=False)
        self.proj = nn.Linear(n_embd, n_embd, bias=False)

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.c_attn(x).split(C, dim=2)
        q = q.view(B, T, self.n_head, self.head_size).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_size).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_size).transpose(1, 2)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(out)


class FeedForward(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.GELU(),
            nn.Linear(4 * n_embd, n_embd),
        )

    def forward(self, x):
        return self.net(x)


class Block(nn.Module):
    def __init__(self, n_embd, n_head):
        super().__init__()
        self.sa = MultiHeadAttention(n_embd, n_head)
        self.ffwd = FeedForward(n_embd)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x


class TinyGPTForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = TinyGPTConfig
    main_input_name = "input_ids"
    _no_split_modules = ["Block"]
    _tied_weights_keys = {"lm_head.weight": "token_embedding_table.weight"}

    def __init__(self, config: TinyGPTConfig):
        super().__init__(config)
        self.token_embedding_table = nn.Embedding(config.vocab_size, config.n_embd)
        self.position_embedding_table = nn.Embedding(config.block_size, config.n_embd)
        self.blocks = nn.Sequential(
            *[Block(config.n_embd, config.n_head) for _ in range(config.n_layer)]
        )
        self.ln_f = nn.LayerNorm(config.n_embd)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        # runs custom _init_weights below on every submodule, then ties
        # lm_head.weight <-> token_embedding_table.weight via
        # get_input_embeddings()/get_output_embeddings() + _tied_weights_keys
        # (same net effect as the manual assignment in the training script,
        # but done the way HF's save/load machinery expects)
        self.post_init()

        # this model has no KV-cache implementation: force generate() to
        # always recompute the full (truncated) sequence each step
        self.generation_config.use_cache = False

    # kept identical to the training script; only used if you ever call
    # .init_weights() on a *fresh* (non-loaded) model
    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)

    # ── HF plumbing for tied weights / embeddings ──
    def get_input_embeddings(self):
        return self.token_embedding_table

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

    def get_output_embeddings(self):
        return self.lm_head

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

    # ── forward pass, same math as the training script ──
    def forward(self, input_ids=None, labels=None, attention_mask=None, **kwargs):
        B, T = input_ids.shape
        device = input_ids.device
        tok_emb = self.token_embedding_table(input_ids)
        pos_emb = self.position_embedding_table(torch.arange(T, device=device))
        x = tok_emb + pos_emb
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                labels.view(-1),
            )

        return CausalLMOutput(loss=loss, logits=logits)

    # ── generation support (no KV cache, so we just resend the
    #    truncated running sequence every step — matches the
    #    behaviour of the model.generate() used at training time) ──
    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        if input_ids.shape[1] > self.config.block_size:
            input_ids = input_ids[:, -self.config.block_size :]
        return {"input_ids": input_ids}