"""
Teensy Language Model

A skinny-deep decoder-only transformer for small-scale language modeling.
Architecture modified from NanoGPT (Andrej Karpathy) and the GPT family.

Copyright (c) 2025 Pankaj Doharey
"""

import math
import inspect
from dataclasses import dataclass

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


class TeensyLayerNorm(nn.Module):
    """Layer normalization with optional bias."""

    def __init__(self, features, use_bias):
        super().__init__()
        self.gain = nn.Parameter(torch.ones(features))
        self.bias = nn.Parameter(torch.zeros(features)) if use_bias else None

    def forward(self, x):
        return F.layer_norm(x, self.gain.shape, self.gain, self.bias, 1e-5)


class TeensySelfAttention(nn.Module):
    """Causal multi-head self-attention with optional Flash Attention."""

    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        head_dim = config.n_embd // config.n_head

        self.qkv_proj = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        self.out_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        self.attn_drop = nn.Dropout(config.dropout)
        self.out_drop = nn.Dropout(config.dropout)

        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.head_dim = head_dim
        self.dropout = config.dropout

        self.use_flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
        if not self.use_flash:
            self.register_buffer(
                "mask",
                torch.tril(torch.ones(config.block_size, config.block_size))
                    .view(1, 1, config.block_size, config.block_size)
            )

    def forward(self, x):
        B, T, C = x.size()

        q, k, v = self.qkv_proj(x).split(self.n_embd, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        if self.use_flash:
            y = torch.nn.functional.scaled_dot_product_attention(
                q, k, v,
                attn_mask=None,
                dropout_p=self.dropout if self.training else 0.0,
                is_causal=True,
            )
        else:
            scores = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
            scores = scores.masked_fill(self.mask[:, :, :T, :T] == 0, float('-inf'))
            weights = F.softmax(scores, dim=-1)
            weights = self.attn_drop(weights)
            y = weights @ v

        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.out_drop(self.out_proj(y))


class TeensyFeedForward(nn.Module):
    """Position-wise feed-forward network with GELU activation."""

    def __init__(self, config):
        super().__init__()
        self.up_proj = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
        self.act = nn.GELU()
        self.down_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        return self.dropout(self.down_proj(self.act(self.up_proj(x))))


class TeensyBlock(nn.Module):
    """Pre-norm transformer block: attention + feed-forward with residuals."""

    def __init__(self, config):
        super().__init__()
        self.attn_norm = TeensyLayerNorm(config.n_embd, config.bias)
        self.attn = TeensySelfAttention(config)
        self.ffn_norm = TeensyLayerNorm(config.n_embd, config.bias)
        self.ffn = TeensyFeedForward(config)

    def forward(self, x):
        x = x + self.attn(self.attn_norm(x))
        x = x + self.ffn(self.ffn_norm(x))
        return x


@dataclass
class TeensyConfig:
    block_size: int = 1024
    vocab_size: int = 50304
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768
    dropout: float = 0.0
    bias: bool = True


class TeensyLM(nn.Module):
    """Teensy decoder-only language model."""

    def __init__(self, config):
        super().__init__()
        assert config.vocab_size is not None
        assert config.block_size is not None
        self.config = config

        self.token_emb = nn.Embedding(config.vocab_size, config.n_embd)
        self.pos_emb = nn.Embedding(config.block_size, config.n_embd)
        self.emb_drop = nn.Dropout(config.dropout)
        self.layers = nn.ModuleList([TeensyBlock(config) for _ in range(config.n_layer)])
        self.final_norm = TeensyLayerNorm(config.n_embd, config.bias)
        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        # Weight tying between token embeddings and output head.
        self.token_emb.weight = self.head.weight

        self.apply(self._init_weights)
        self._init_residual_projections()

        print(f"number of parameters: {self.get_num_params() / 1e6:.2f}M")

    def get_num_params(self, non_embedding=True):
        n = sum(p.numel() for p in self.parameters())
        if non_embedding:
            n -= self.pos_emb.weight.numel()
        return n

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            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=0.02)

    def _init_residual_projections(self):
        for name, param in self.named_parameters():
            if name.endswith(('attn.out_proj.weight', 'ffn.down_proj.weight')):
                torch.nn.init.normal_(param, mean=0.0, std=0.02 / math.sqrt(2 * self.config.n_layer))

    def forward(self, idx, targets=None):
        B, T = idx.size()
        assert T <= self.config.block_size, (
            f"Sequence length {T} exceeds block size {self.config.block_size}"
        )

        positions = torch.arange(T, dtype=torch.long, device=idx.device)
        x = self.emb_drop(self.token_emb(idx) + self.pos_emb(positions))
        for layer in self.layers:
            x = layer(x)
        x = self.final_norm(x)

        if targets is not None:
            logits = self.head(x)
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1),
                ignore_index=-1,
            )
        else:
            logits = self.head(x[:, [-1], :])
            loss = None

        return logits, loss

    def crop_block_size(self, block_size):
        assert block_size <= self.config.block_size
        self.config.block_size = block_size
        self.pos_emb.weight = nn.Parameter(self.pos_emb.weight[:block_size])
        for layer in self.layers:
            if hasattr(layer.attn, 'mask'):
                layer.attn.mask = layer.attn.mask[:, :, :block_size, :block_size]

    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
        params = {n: p for n, p in self.named_parameters() if p.requires_grad}
        decay = [p for p in params.values() if p.dim() >= 2]
        no_decay = [p for p in params.values() if p.dim() < 2]

        groups = [
            {'params': decay, 'weight_decay': weight_decay},
            {'params': no_decay, 'weight_decay': 0.0},
        ]

        print(f"decayed tensors: {len(decay)}, params: {sum(p.numel() for p in decay):,}")
        print(f"non-decayed tensors: {len(no_decay)}, params: {sum(p.numel() for p in no_decay):,}")

        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
        use_fused = fused_available and device_type == 'cuda'
        optimizer = torch.optim.AdamW(groups, lr=learning_rate, betas=betas, **({'fused': True} if use_fused else {}))
        print(f"using fused AdamW: {use_fused}")
        return optimizer

    def estimate_mfu(self, fwdbwd_per_iter, dt):
        """Estimate model FLOPs utilization against an A100 bfloat16 peak."""
        N = self.get_num_params()
        cfg = self.config
        L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd // cfg.n_head, cfg.block_size
        flops_per_token = 6 * N + 12 * L * H * Q * T
        flops_per_iter = flops_per_token * T * fwdbwd_per_iter
        flops_per_sec = flops_per_iter / dt
        a100_peak = 312e12
        return flops_per_sec / a100_peak

    @staticmethod
    def _apply_top_p(logits, top_p):
        """Nucleus (top-p) sampling: zero out logits outside the smallest nucleus.

        logits: (B, V) unnormalized logits.
        top_p: cumulative probability threshold in (0, 1].
        """
        if top_p is None or top_p <= 0.0 or top_p >= 1.0:
            return logits
        probs = F.softmax(logits, dim=-1)
        sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1)
        cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
        # Keep tokens whose cumulative probability is <= top_p; also keep the first token over the threshold.
        sorted_indices_to_remove = cumulative_probs > top_p
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = False
        indices_to_remove = sorted_indices_to_remove.scatter(-1, sorted_indices, sorted_indices_to_remove)
        logits = logits.masked_fill(indices_to_remove, -float('inf'))
        return logits

    @staticmethod
    def _sample_next(logits, temperature=1.0, top_k=None, top_p=None):
        """Sample one next-token id from last-position logits.

        temperature <= 0 switches to greedy argmax decoding. This makes
        deterministic generation explicit and avoids divide-by-zero/NaN output.
        """
        if temperature is None or temperature <= 0.0:
            return torch.argmax(logits, dim=-1, keepdim=True)

        logits = logits / temperature

        if top_k is not None and top_k > 0:
            k = min(top_k, logits.size(-1))
            top_vals, _ = torch.topk(logits, k)
            logits = logits.masked_fill(logits < top_vals[:, [-1]], -float('inf'))

        logits = TeensyLM._apply_top_p(logits, top_p)

        probs = F.softmax(logits, dim=-1)
        return torch.multinomial(probs, num_samples=1)

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None, top_p=None, eos_token_id=None):
        """Autoregressively generate tokens from a conditioning sequence.

        If eos_token_id is provided, generation stops when that token is emitted
        and the EOS token is trimmed from the returned sequence.
        """
        for _ in range(max_new_tokens):
            ctx = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
            logits, _ = self(ctx)
            next_token = self._sample_next(logits[:, -1, :], temperature, top_k, top_p)
            idx = torch.cat((idx, next_token), dim=1)

            if eos_token_id is not None and next_token.item() == eos_token_id:
                break

        if eos_token_id is not None and idx[0, -1].item() == eos_token_id:
            idx = idx[:, :-1]
        return idx

    @torch.no_grad()
    def generate_stream(self, idx, max_new_tokens, temperature=1.0, top_k=None, top_p=None, eos_token_id=None):
        """
        Autoregressively generate tokens and yield each token id as it is produced.
        The caller can decode and print tokens incrementally for a streaming UX.

        If eos_token_id is provided, generation stops when that token is emitted.
        """
        for _ in range(max_new_tokens):
            ctx = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
            logits, _ = self(ctx)
            next_token = self._sample_next(logits[:, -1, :], temperature, top_k, top_p)
            idx = torch.cat((idx, next_token), dim=1)
            token_id = next_token.item()

            if eos_token_id is not None and token_id == eos_token_id:
                break
            yield token_id


def adapt_nanogpt_weights(state_dict):
    """
    Rename weights from the original NanoGPT naming scheme to the Teensy scheme.
    This lets checkpoints produced by the original training code load into the
    refactored Teensy model without retraining.
    """
    mapping = {
        'transformer.wte.weight': 'token_emb.weight',
        'transformer.wpe.weight': 'pos_emb.weight',
        'transformer.ln_f.weight': 'final_norm.gain',
        'transformer.ln_f.bias': 'final_norm.bias',
        'lm_head.weight': 'head.weight',
    }

    layer_mappings = {
        'ln_1.weight': 'attn_norm.gain',
        'ln_1.bias': 'attn_norm.bias',
        'attn.c_attn.weight': 'attn.qkv_proj.weight',
        'attn.c_attn.bias': 'attn.qkv_proj.bias',
        'attn.c_proj.weight': 'attn.out_proj.weight',
        'attn.c_proj.bias': 'attn.out_proj.bias',
        'ln_2.weight': 'ffn_norm.gain',
        'ln_2.bias': 'ffn_norm.bias',
        'mlp.c_fc.weight': 'ffn.up_proj.weight',
        'mlp.c_fc.bias': 'ffn.up_proj.bias',
        'mlp.c_proj.weight': 'ffn.down_proj.weight',
        'mlp.c_proj.bias': 'ffn.down_proj.bias',
    }

    adapted = {}
    for old_key, tensor in state_dict.items():
        if old_key in mapping:
            new_key = mapping[old_key]
        elif old_key.startswith('transformer.h.'):
            parts = old_key.split('.')
            layer_idx = parts[2]
            sub_key = '.'.join(parts[3:])
            if sub_key in layer_mappings:
                new_key = f'layers.{layer_idx}.{layer_mappings[sub_key]}'
            else:
                new_key = old_key
        else:
            new_key = old_key
        adapted[new_key] = tensor
    return adapted