"""Compact decoder-only Transformer for the Part 2 chat experiment."""
from dataclasses import asdict, dataclass
import math

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


@dataclass
class ModelConfig:
    vocab_size: int = 8192
    d_model: int = 192
    n_layers: int = 6
    n_heads: int = 6
    n_kv_heads: int = 2
    d_ff: int = 512
    max_seq_len: int = 512
    dropout: float = 0.0


def parameter_count(config: ModelConfig) -> int:
    # tied token/output embedding; Q and output projections; grouped K/V.
    d, h, kv, f, v = config.d_model, config.n_heads, config.n_kv_heads, config.d_ff, config.vocab_size
    attn = d * d + 2 * d * (d * kv // h) + d * d
    # SwiGLU has gate/up/down projections. RMSNorm has one scale per norm.
    layer = attn + 3 * d * f + 2 * d
    return v * d + config.n_layers * layer + d


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

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


class Attention(nn.Module):
    def __init__(self, c: ModelConfig):
        super().__init__()
        assert c.d_model % c.n_heads == 0 and c.n_heads % c.n_kv_heads == 0
        self.head_dim = c.d_model // c.n_heads
        self.n_heads, self.n_kv_heads = c.n_heads, c.n_kv_heads
        self.q = nn.Linear(c.d_model, c.d_model, bias=False)
        kv_dim = c.n_kv_heads * self.head_dim
        self.k = nn.Linear(c.d_model, kv_dim, bias=False)
        self.v = nn.Linear(c.d_model, kv_dim, bias=False)
        self.out = nn.Linear(c.d_model, c.d_model, bias=False)
        self.dropout = c.dropout

    def forward(self, x):
        b, t, _ = x.shape
        q = self.q(x).view(b, t, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2)
        if self.n_kv_heads != self.n_heads:
            repeat = self.n_heads // self.n_kv_heads
            k, v = k.repeat_interleave(repeat, dim=1), v.repeat_interleave(repeat, dim=1)
        # RoPE, with interleaved even/odd pairs.
        half = self.head_dim // 2
        freq = torch.arange(half, device=x.device, dtype=torch.float32)
        inv_freq = 1.0 / (10000 ** (freq / half))
        pos = torch.arange(t, device=x.device, dtype=torch.float32)
        angle = torch.outer(pos, inv_freq)
        cos, sin = angle.cos()[None, None], angle.sin()[None, None]

        def rotate(z):
            z1, z2 = z[..., :half], z[..., half:2 * half]
            return torch.cat((z1 * cos - z2 * sin, z1 * sin + z2 * cos, z[..., 2 * half:]), dim=-1).to(z.dtype)

        q, k = rotate(q), rotate(k)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True,
                                           dropout_p=self.dropout if self.training else 0.0)
        return self.out(y.transpose(1, 2).contiguous().view(b, t, -1))


class Block(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.attn_norm, self.ffn_norm = RMSNorm(c.d_model), RMSNorm(c.d_model)
        self.attn = Attention(c)
        self.gate = nn.Linear(c.d_model, c.d_ff, bias=False)
        self.up = nn.Linear(c.d_model, c.d_ff, bias=False)
        self.down = nn.Linear(c.d_ff, c.d_model, bias=False)
        self.dropout = nn.Dropout(c.dropout)

    def forward(self, x):
        x = x + self.dropout(self.attn(self.attn_norm(x)))
        x = x + self.dropout(self.down(F.silu(self.gate(self.ffn_norm(x))) * self.up(self.ffn_norm(x))))
        return x


class TinyChatLM(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layers)])
        self.norm = RMSNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.lm_head.weight = self.token_embedding.weight
        self.apply(self._init_weights)
        for name, p in self.named_parameters():
            if name.endswith("out.weight") or name.endswith("down.weight"):
                nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layers))

    @staticmethod
    def _init_weights(module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, ids, targets=None, loss_mask=None):
        x = self.token_embedding(ids)
        for block in self.blocks:
            x = block(x)
        logits = self.lm_head(self.norm(x))
        loss = None
        if targets is not None:
            per_token = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), reduction="none")
            if loss_mask is not None:
                weights = loss_mask.reshape(-1).to(per_token.dtype)
                loss = (per_token * weights).sum() / weights.sum().clamp_min(1)
            else:
                loss = per_token.mean()
        return logits, loss

    @torch.no_grad()
    def generate(self, ids, max_new_tokens=100, temperature=0.8, top_k=40, eos_id=None):
        self.eval()
        for _ in range(max_new_tokens):
            ctx = ids[:, -self.config.max_seq_len:]
            logits, _ = self(ctx)
            logits = logits[:, -1, :] / max(temperature, 1e-5)
            if top_k:
                cutoff = torch.topk(logits, min(top_k, logits.size(-1))).values[:, -1:]
                logits = logits.masked_fill(logits < cutoff, -float("inf"))
            nxt = torch.multinomial(F.softmax(logits, dim=-1), 1)
            ids = torch.cat((ids, nxt), dim=1)
            if eos_id is not None and (nxt == eos_id).all():
                break
        return ids


def config_dict(config):
    return asdict(config)