"""Training loop and evaluation harness for the lesson 35 GPT model.

Implements: batch construction with input/target shift by one, cross entropy
loss in `calc_loss_batch`, held-out evaluation in `evaluate_model`, a qualitative
generation probe in `generate_and_print_sample`, AdamW with a decay/no-decay
split, a linear-warmup-plus-cosine learning rate schedule, gradient norm
clipping, and a JSONL log of per step loss in `outputs/losses.jsonl`.

The demo trains a tiny model on synthetic byte-level tokens for a small number
of steps, writes the JSONL log, and prints eval losses and generated samples
at the probe points. End to end runs in well under a minute on CPU.

Run: python3 code/main.py
"""

from __future__ import annotations

import json
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Iterator

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

HERE = Path(__file__).resolve().parent
OUTPUTS = HERE.parent / "outputs"
OUTPUTS.mkdir(parents=True, exist_ok=True)
LOG_PATH = OUTPUTS / "losses.jsonl"


@dataclass
class TrainConfig:
    """Training and evaluation hyperparameters for the demo run."""

    batch_size: int = 4
    context_length: int = 32
    num_steps: int = 80
    eval_every: int = 20
    eval_batches: int = 4
    max_lr: float = 3e-3
    min_lr: float = 3e-4
    warmup_steps: int = 10
    weight_decay: float = 0.01
    grad_clip: float = 1.0
    sample_max_new_tokens: int = 16
    seed: int = 0


@dataclass
class ModelConfig:
    vocab_size: int = 256
    context_length: int = 32
    d_model: int = 64
    num_heads: int = 4
    num_layers: int = 2
    mlp_expansion: int = 4
    dropout: float = 0.1
    use_bias: bool = True
    weight_tying: bool = True


class LayerNorm(nn.Module):
    def __init__(self, d_model: int, eps: float = 1e-5) -> None:
        super().__init__()
        self.eps = eps
        self.scale = nn.Parameter(torch.ones(d_model))
        self.shift = nn.Parameter(torch.zeros(d_model))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        return self.scale * (x - mean) / torch.sqrt(var + self.eps) + self.shift


class MultiHeadAttention(nn.Module):
    def __init__(self, cfg: ModelConfig) -> None:
        super().__init__()
        if cfg.d_model % cfg.num_heads != 0:
            raise ValueError("d_model must be divisible by num_heads")
        self.d_model = cfg.d_model
        self.num_heads = cfg.num_heads
        self.head_dim = cfg.d_model // cfg.num_heads
        self.context_length = cfg.context_length
        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=cfg.use_bias)
        self.out_proj = nn.Linear(cfg.d_model, cfg.d_model, bias=cfg.use_bias)
        self.attn_dropout = nn.Dropout(cfg.dropout)
        self.resid_dropout = nn.Dropout(cfg.dropout)
        mask = torch.triu(
            torch.ones(cfg.context_length, cfg.context_length, dtype=torch.bool),
            diagonal=1,
        )
        self.register_buffer("causal_mask", mask, persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch, seq, dim = x.shape
        qkv = self.qkv(x)
        q, k, v = qkv.split(self.d_model, dim=-1)
        q = q.view(batch, seq, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch, seq, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch, seq, self.num_heads, self.head_dim).transpose(1, 2)
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.head_dim)
        scores = scores.masked_fill(self.causal_mask[:seq, :seq], float("-inf"))
        attn = F.softmax(scores, dim=-1)
        attn = self.attn_dropout(attn)
        out = (attn @ v).transpose(1, 2).contiguous().view(batch, seq, dim)
        return self.resid_dropout(self.out_proj(out))


class FeedForward(nn.Module):
    def __init__(self, cfg: ModelConfig) -> None:
        super().__init__()
        hidden = cfg.mlp_expansion * cfg.d_model
        self.fc1 = nn.Linear(cfg.d_model, hidden, bias=cfg.use_bias)
        self.act = nn.GELU(approximate="tanh")
        self.fc2 = nn.Linear(hidden, cfg.d_model, bias=cfg.use_bias)
        self.dropout = nn.Dropout(cfg.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.dropout(self.fc2(self.act(self.fc1(x))))


class TransformerBlock(nn.Module):
    def __init__(self, cfg: ModelConfig) -> None:
        super().__init__()
        self.ln1 = LayerNorm(cfg.d_model)
        self.attn = MultiHeadAttention(cfg)
        self.ln2 = LayerNorm(cfg.d_model)
        self.mlp = FeedForward(cfg)

    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 GPTModel(nn.Module):
    def __init__(self, cfg: ModelConfig) -> None:
        super().__init__()
        self.cfg = cfg
        self.tok_embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.pos_embed = nn.Embedding(cfg.context_length, cfg.d_model)
        self.embed_dropout = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg.num_layers)])
        self.final_ln = LayerNorm(cfg.d_model)
        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        if cfg.weight_tying:
            self.lm_head.weight = self.tok_embed.weight
        self.register_buffer(
            "position_ids",
            torch.arange(cfg.context_length, dtype=torch.long),
            persistent=False,
        )
        self.apply(self._init_weights)
        scale = 1.0 / math.sqrt(2 * cfg.num_layers)
        for block in self.blocks:
            block.attn.out_proj.weight.data.mul_(scale)
            block.mlp.fc2.weight.data.mul_(scale)

    @staticmethod
    def _init_weights(module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
        batch, seq = tokens.shape
        tok = self.tok_embed(tokens)
        pos = self.pos_embed(self.position_ids[:seq])
        x = self.embed_dropout(tok + pos)
        for block in self.blocks:
            x = block(x)
        return self.lm_head(self.final_ln(x))


def make_batches(
    token_ids: torch.Tensor,
    batch_size: int,
    context_length: int,
    seed: int = 0,
) -> Iterator[tuple[torch.Tensor, torch.Tensor]]:
    """Yield (input, target) batches where target is input shifted by one position.

    Sampling is uniform random over valid start positions. With a fixed seed the
    sequence of batches is reproducible across runs.
    """
    if token_ids.dim() != 1:
        raise ValueError("token_ids must be a 1D tensor")
    if token_ids.numel() < context_length + 1:
        raise ValueError("token_ids too short for the requested context_length")

    generator = torch.Generator().manual_seed(seed)
    max_start = token_ids.numel() - context_length - 1

    while True:
        starts = torch.randint(0, max_start + 1, (batch_size,), generator=generator)
        inputs = torch.stack([token_ids[s : s + context_length] for s in starts.tolist()])
        targets = torch.stack(
            [token_ids[s + 1 : s + 1 + context_length] for s in starts.tolist()]
        )
        yield inputs, targets


def calc_loss_batch(
    model: GPTModel,
    inputs: torch.Tensor,
    targets: torch.Tensor,
) -> torch.Tensor:
    """Forward, flatten across batch and time, return scalar cross entropy."""
    logits = model(inputs)
    return F.cross_entropy(
        logits.reshape(-1, logits.size(-1)),
        targets.reshape(-1),
    )


@torch.no_grad()
def evaluate_model(
    model: GPTModel,
    val_loader: Iterator[tuple[torch.Tensor, torch.Tensor]],
    max_batches: int,
) -> float:
    """Mean cross entropy over `max_batches` validation batches; no grad, no dropout."""
    was_training = model.training
    model.eval()
    total = 0.0
    count = 0
    for inputs, targets in val_loader:
        if count >= max_batches:
            break
        loss = calc_loss_batch(model, inputs, targets)
        total += float(loss.item())
        count += 1
    if was_training:
        model.train()
    return total / max(count, 1)


@torch.no_grad()
def generate_and_print_sample(
    model: GPTModel,
    prompt: torch.Tensor,
    max_new_tokens: int,
    temperature: float = 1.0,
    top_k: int = 40,
    seed: int = 0,
) -> list[int]:
    """Print a short generated continuation from a fixed prompt and return the tokens."""
    sample_gen = torch.Generator(device=prompt.device).manual_seed(seed)
    was_training = model.training
    model.eval()
    tokens = prompt.clone()
    for _ in range(max_new_tokens):
        window = tokens[:, -model.cfg.context_length :]
        logits = model(window)
        next_logits = logits[:, -1, :] / temperature
        if top_k > 0:
            top_k_eff = min(top_k, next_logits.size(-1))
            values, _ = torch.topk(next_logits, top_k_eff, dim=-1)
            threshold = values[..., -1:]
            next_logits = torch.where(
                next_logits < threshold, torch.full_like(next_logits, float("-inf")), next_logits
            )
        probs = F.softmax(next_logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1, generator=sample_gen)
        tokens = torch.cat([tokens, next_token], dim=1)
    if was_training:
        model.train()
    seq = tokens.tolist()[0]
    print(f"  sample tokens          : {seq}")
    return seq


def build_param_groups(model: nn.Module, weight_decay: float) -> list[dict]:
    """Split parameters: matrix-shaped tensors get decay; scale/bias/embedding biases do not."""
    decay: list[nn.Parameter] = []
    no_decay: list[nn.Parameter] = []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if param.dim() < 2 or name.endswith(".bias") or name.endswith(".shift") or name.endswith(".scale"):
            no_decay.append(param)
        else:
            decay.append(param)
    return [
        {"params": decay, "weight_decay": weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]


def cosine_with_warmup(
    step: int,
    warmup_steps: int,
    total_steps: int,
    max_lr: float,
    min_lr: float,
) -> float:
    """Linear warmup then cosine decay to min_lr over the remaining steps."""
    if step < warmup_steps:
        return max_lr * (step + 1) / max(warmup_steps, 1)
    progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1)
    progress = min(max(progress, 0.0), 1.0)
    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
    return min_lr + (max_lr - min_lr) * cosine


def train(
    model: GPTModel,
    train_tokens: torch.Tensor,
    val_tokens: torch.Tensor,
    cfg: TrainConfig,
    prompt: torch.Tensor,
    log_path: Path = LOG_PATH,
) -> list[dict]:
    """Run the training loop, persist losses.jsonl, return the in-memory log."""
    torch.manual_seed(cfg.seed)
    optimizer = torch.optim.AdamW(
        build_param_groups(model, cfg.weight_decay),
        lr=cfg.max_lr,
        betas=(0.9, 0.95),
    )

    train_loader = make_batches(train_tokens, cfg.batch_size, cfg.context_length, seed=cfg.seed)

    if log_path.exists():
        log_path.unlink()

    records: list[dict] = []
    model.train()
    for step in range(cfg.num_steps):
        lr = cosine_with_warmup(step, cfg.warmup_steps, cfg.num_steps, cfg.max_lr, cfg.min_lr)
        for group in optimizer.param_groups:
            group["lr"] = lr

        inputs, targets = next(train_loader)
        optimizer.zero_grad(set_to_none=True)
        loss = calc_loss_batch(model, inputs, targets)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=cfg.grad_clip)
        optimizer.step()

        record = {"step": step, "train_loss": float(loss.item()), "lr": lr}

        if (step + 1) % cfg.eval_every == 0 or step == cfg.num_steps - 1:
            val_loader = make_batches(
                val_tokens, cfg.batch_size, cfg.context_length, seed=cfg.seed + 1
            )
            val_loss = evaluate_model(model, val_loader, cfg.eval_batches)
            record["val_loss"] = val_loss
            print(
                f"step {step:4d} | lr {lr:.5f} | train_loss {loss.item():.4f} | val_loss {val_loss:.4f}"
            )
            generate_and_print_sample(
                model, prompt, cfg.sample_max_new_tokens, temperature=0.8, top_k=20, seed=step
            )

        records.append(record)
        with log_path.open("a", encoding="utf-8") as fh:
            fh.write(json.dumps(record) + "\n")

    return records


def _synthetic_byte_tokens(length: int, vocab_size: int, seed: int) -> torch.Tensor:
    """Deterministic synthetic tokens.

    Bytes drawn from a small repeating pattern so the model has structure to learn
    in a handful of steps; the eval loss should drop visibly during the demo.
    """
    rng = torch.Generator().manual_seed(seed)
    base = torch.randint(0, vocab_size, (32,), generator=rng)
    repeats = (length + base.numel() - 1) // base.numel()
    tokens = base.repeat(repeats)[:length]
    noise = torch.randint(0, vocab_size, (length,), generator=rng)
    mask = (torch.rand(length, generator=rng) < 0.1)
    tokens = torch.where(mask, noise, tokens)
    return tokens.to(dtype=torch.long)


def demo() -> None:
    torch.manual_seed(0)
    cfg = TrainConfig()
    mcfg = ModelConfig(
        vocab_size=256,
        context_length=cfg.context_length,
        d_model=64,
        num_heads=4,
        num_layers=2,
        dropout=0.0,
    )

    train_tokens = _synthetic_byte_tokens(length=4096, vocab_size=mcfg.vocab_size, seed=1)
    val_tokens = _synthetic_byte_tokens(length=1024, vocab_size=mcfg.vocab_size, seed=2)

    print(f"train tokens   : {train_tokens.numel():,}")
    print(f"val tokens     : {val_tokens.numel():,}")
    print(f"model params   : {sum(p.numel() for p in GPTModel(mcfg).parameters()):,} (untied count)")

    model = GPTModel(mcfg)
    prompt = torch.tensor([[7, 11, 13, 17]], dtype=torch.long)

    print("\nTraining run:")
    records = train(model, train_tokens, val_tokens, cfg, prompt)

    print("\nFinal log records (last 3):")
    for record in records[-3:]:
        print(" ", record)
    print(f"\nWrote losses to {LOG_PATH}")

    first_loss = records[0]["train_loss"]
    last_loss = records[-1]["train_loss"]
    print(f"First step train_loss  : {first_loss:.4f}")
    print(f"Last step train_loss   : {last_loss:.4f}")
    assert last_loss < first_loss, "training loss should decrease across the demo run"
    print("Training loop check passed.")


if __name__ == "__main__":
    demo()