"""
Derived from Andrej Karpathy's nanochat project.

MIT License

Copyright (c) 2025 Andrej Karpathy

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass
import math

import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass(frozen=True)
class GPTConfig:
    block_size: int
    vocab_size: int
    n_layer: int
    n_head: int
    n_embd: int
    dropout: float

    def to_dict(self) -> dict[str, int | float]:
        return asdict(self)


def rms_norm(x: torch.Tensor) -> torch.Tensor:
    return F.rms_norm(x, (x.size(-1),))


class Linear(nn.Linear):
    """Bias-free linear layer matching the nanochat architectural style."""

    def __init__(self, in_features: int, out_features: int):
        super().__init__(in_features, out_features, bias=False)


def apply_rotary_emb(
    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
) -> torch.Tensor:
    half = x.shape[-1] // 2
    x1, x2 = x[..., :half], x[..., half:]
    return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1)


class CausalSelfAttention(nn.Module):
    def __init__(self, config: GPTConfig):
        super().__init__()
        if config.n_embd % config.n_head != 0:
            raise ValueError("n_embd must be divisible by n_head")
        self.n_head = config.n_head
        self.head_dim = config.n_embd // config.n_head
        if self.head_dim % 2 != 0:
            raise ValueError("rotary attention requires an even head dimension")
        self.dropout_p = config.dropout
        self.c_q = Linear(config.n_embd, config.n_embd)
        self.c_k = Linear(config.n_embd, config.n_embd)
        self.c_v = Linear(config.n_embd, config.n_embd)
        self.c_proj = Linear(config.n_embd, config.n_embd)
        self.resid_dropout = nn.Dropout(config.dropout)
        mask = torch.ones(config.block_size, config.block_size, dtype=torch.bool).tril()
        self.register_buffer(
            "causal_mask",
            mask.view(1, 1, config.block_size, config.block_size),
            persistent=False,
        )

    def set_dropout(self, p: float) -> None:
        self.dropout_p = p
        self.resid_dropout.p = p

    def forward(
        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
    ) -> torch.Tensor:
        batch, seq_len, channels = x.shape
        q = self.c_q(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)
        k = self.c_k(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)
        v = self.c_v(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)

        q = rms_norm(apply_rotary_emb(q, cos, sin)) * 1.2
        k = rms_norm(apply_rotary_emb(k, cos, sin)) * 1.2

        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        mask = self.causal_mask[:, :, :seq_len, :seq_len]
        att = att.masked_fill(~mask, float("-inf"))
        att = F.softmax(att.float(), dim=-1).to(dtype=x.dtype)
        att = F.dropout(att, p=self.dropout_p, training=self.training)
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(batch, seq_len, channels)
        return self.resid_dropout(self.c_proj(y))


class MLP(nn.Module):
    def __init__(self, config: GPTConfig):
        super().__init__()
        self.c_fc = Linear(config.n_embd, 4 * config.n_embd)
        self.c_proj = Linear(4 * config.n_embd, config.n_embd)
        self.dropout = nn.Dropout(config.dropout)

    def set_dropout(self, p: float) -> None:
        self.dropout.p = p

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.c_fc(x)
        x = F.relu(x).square()
        x = self.c_proj(x)
        return self.dropout(x)


class Block(nn.Module):
    def __init__(self, config: GPTConfig):
        super().__init__()
        self.attn = CausalSelfAttention(config)
        self.mlp = MLP(config)

    def set_dropout(self, p: float) -> None:
        self.attn.set_dropout(p)
        self.mlp.set_dropout(p)

    def forward(
        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
    ) -> torch.Tensor:
        x = x + self.attn(rms_norm(x), cos, sin)
        x = x + self.mlp(rms_norm(x))
        return x


class DropoutGPT(nn.Module):
    """Minimal nanochat-style causal Transformer with dynamic dropout control."""

    def __init__(self, config: GPTConfig):
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
        self.embedding_dropout = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList(Block(config) for _ in range(config.n_layer))
        self.lm_head = Linear(config.n_embd, config.vocab_size)
        self.register_buffer(
            "cos",
            torch.empty(1, 1, config.block_size, config.n_embd // config.n_head // 2),
            persistent=False,
        )
        self.register_buffer(
            "sin",
            torch.empty(1, 1, config.block_size, config.n_embd // config.n_head // 2),
            persistent=False,
        )
        self._init_weights()
        self._init_rotary()

    @torch.no_grad()
    def _init_weights(self) -> None:
        nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.8)
        nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.001)
        scale = (3.0 ** 0.5) * (self.config.n_embd ** -0.5)
        for block in self.blocks:
            nn.init.uniform_(block.attn.c_q.weight, -scale, scale)
            nn.init.uniform_(block.attn.c_k.weight, -scale, scale)
            nn.init.uniform_(block.attn.c_v.weight, -scale, scale)
            nn.init.zeros_(block.attn.c_proj.weight)
            nn.init.uniform_(block.mlp.c_fc.weight, -0.4 * scale, 0.4 * scale)
            nn.init.zeros_(block.mlp.c_proj.weight)

    @torch.no_grad()
    def _init_rotary(self) -> None:
        head_dim = self.config.n_embd // self.config.n_head
        channel_range = torch.arange(
            0,
            head_dim,
            2,
            dtype=torch.float32,
            device=self.cos.device,
        )
        inv_freq = 1.0 / (100000 ** (channel_range / head_dim))
        positions = torch.arange(
            self.config.block_size,
            dtype=torch.float32,
            device=self.cos.device,
        )
        freqs = torch.outer(positions, inv_freq)
        self.cos.copy_(freqs.cos().view(1, 1, self.config.block_size, head_dim // 2))
        self.sin.copy_(freqs.sin().view(1, 1, self.config.block_size, head_dim // 2))

    def set_dropout(self, p: float) -> None:
        if not 0.0 <= p < 1.0:
            raise ValueError("dropout must satisfy 0 <= p < 1")
        self.embedding_dropout.p = p
        for block in self.blocks:
            block.set_dropout(p)

    def num_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters())

    def forward(
        self, idx: torch.Tensor, targets: torch.Tensor | None = None
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        _, seq_len = idx.shape
        if seq_len > self.config.block_size:
            raise ValueError("sequence length exceeds block_size")
        cos = self.cos[:, :, :seq_len, :]
        sin = self.sin[:, :, :seq_len, :]
        x = self.embedding_dropout(rms_norm(self.token_embedding(idx)))
        for block in self.blocks:
            x = block(x, cos, sin)
        logits = self.lm_head(rms_norm(x)).float()
        if targets is None:
            return logits
        loss = F.cross_entropy(
            logits.reshape(-1, logits.size(-1)),
            targets.reshape(-1),
        )
        return logits, loss