"""Tiny random Gemma 4 checks before any 12B run. Run as ``python -m training.selftest [--model DIR] --out FILE``.

Classifier-plan §3.5. The tiny checkpoint is written in the real format (``Gemma4UnifiedForConditionalGeneration``
with small image and audio projections, so the text-only loader must map and ignore keys exactly as for the 12B),
with a byte-level tokenizer that has Gemma's special-token ids (pad 0, eos 1, bos 2). With ``--model`` the real
config's attention shapes (head_dim 256, global head_dim 512, KV heads, sliding window) are kept and only the depth,
width and vocabulary shrink; without it everything is small (CPU tests). Nothing here needs the 12B weights.

Checks (exit 1 if any fails):
1. text-only load: zero missing text weights, multimodal weights ignored, hidden states equal the full model's;
2. heads: one ``Linear(hidden -> K_f)`` per nsb-v1 field (the real head sizes), LoRA on every target projection;
3. padding invariance: a row alone and the same row right-padded in a batch give the same logits (both poolings,
   float64 so that only a masking or pooling error can fail it);
4. loss split: one step as one micro-batch and as two give the same gradients;
5. save/load: adapter + heads reload to the same logits;
6. CPU/GPU parity (only with CUDA): the same tiny model gives the same logits on both devices.
"""
from __future__ import annotations

import argparse
import json
import tempfile
from pathlib import Path

VOCAB = 512
SPECIALS = ["<pad>", "<eos>", "<bos>", "<unk>", "<|turn>", "<turn|>", "<|channel>", "<channel|>"]
LORA = dict(lora_r=8, lora_alpha=8, lora_dropout=0.0, head_dropout=0.0)


def tiny_text(real=None, layers: int = 6) -> dict:
    """Text config kwargs: 5 sliding + 1 full attention layer (the 12B pattern repeated once)."""
    kinds = ["sliding_attention"] * (layers - 1) + ["full_attention"]
    if real is None:
        return dict(vocab_size=VOCAB, hidden_size=64, intermediate_size=128, num_hidden_layers=layers,
                    num_attention_heads=4, num_key_value_heads=2, head_dim=16, global_head_dim=32,
                    num_global_key_value_heads=1, attention_k_eq_v=True, sliding_window=16,
                    max_position_embeddings=4096, final_logit_softcapping=30.0, layer_types=kinds,
                    pad_token_id=0, eos_token_id=1, bos_token_id=2)
    # Read the serialised values: per-layer attributes (KV heads, head_dim) refuse plain attribute access.
    d = real.to_dict()
    keep = ("num_attention_heads", "num_key_value_heads", "head_dim", "sliding_window", "attention_k_eq_v",
            "hidden_activation", "rms_norm_eps", "final_logit_softcapping", "rope_parameters")
    out = {k: d[k] for k in keep if d.get(k) is not None}
    layer = next(iter((d.get("per_layer_config") or {}).values()), {})
    if layer.get("head_dim") is not None:
        out["global_head_dim"] = layer["head_dim"]
    if layer.get("num_key_value_heads") is not None:
        out["num_global_key_value_heads"] = layer["num_key_value_heads"]
    heads = out.get("num_attention_heads", 4)
    out.update(vocab_size=VOCAB, hidden_size=heads * 32, intermediate_size=heads * 64, num_hidden_layers=layers,
               layer_types=kinds, max_position_embeddings=4096, pad_token_id=0, eos_token_id=1, bos_token_id=2)
    return out


def make_tokenizer(path: Path):
    """Byte-level tokenizer (every text encodes) with Gemma's special-token ids; saved next to the model."""
    from tokenizers import Tokenizer, decoders, models, pre_tokenizers
    from transformers import PreTrainedTokenizerFast

    alphabet = sorted(pre_tokenizers.ByteLevel.alphabet())
    vocab = {t: i for i, t in enumerate(SPECIALS)}
    for ch in alphabet:
        vocab.setdefault(ch, len(vocab))
    tk = Tokenizer(models.BPE(vocab=vocab, merges=[], unk_token="<unk>"))
    tk.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False)
    tk.decoder = decoders.ByteLevel()
    tok = PreTrainedTokenizerFast(tokenizer_object=tk, pad_token="<pad>", eos_token="<eos>", bos_token="<bos>",
                                  unk_token="<unk>", additional_special_tokens=SPECIALS[4:])
    tok.save_pretrained(str(path))
    return tok


def make_tiny(path: Path, real_model: str | None = None, layers: int = 6, seed: int = 0, multimodal: bool = True):
    """Save a random tiny Gemma 4 checkpoint (full multimodal class) plus tokenizer; returns the full model."""
    import torch
    from transformers import Gemma4UnifiedConfig, Gemma4UnifiedForConditionalGeneration

    real = None
    if real_model:
        from .model import text_config

        real = text_config(real_model)
    text = tiny_text(real, layers)
    vision = dict(patch_size=2, pooling_kernel_size=2, mm_embed_dim=16, mm_posemb_size=8,
                  output_proj_dims=16) if multimodal else None
    audio = dict(audio_embed_dim=8) if multimodal else None
    cfg = Gemma4UnifiedConfig(text_config=text, vision_config=vision, audio_config=audio)
    torch.manual_seed(seed)
    full = Gemma4UnifiedForConditionalGeneration(cfg).float()
    path = Path(path)
    full.save_pretrained(str(path))
    make_tokenizer(path)
    return full


def tiny_cfg(path, pool: str = "last", dtype: str = "float32") -> dict:
    return dict(LORA, model=str(path), dtype=dtype, pool=pool, gradient_checkpointing=True)


def build_tiny(path, device, pool: str = "last", sizes=None, seed: int = 7, dtype: str = "float32"):
    """Tiny heads model with non-zero LoRA B weights, so LoRA A gradients and the adapter are non-trivial."""
    import torch

    from .model import build
    from .taxonomy import Taxonomy

    sizes = sizes or Taxonomy.default().sizes()
    model = build(tiny_cfg(path, pool, dtype), sizes, device, train=True)
    g = torch.Generator(device="cpu").manual_seed(seed)
    with torch.no_grad():
        for n, p in model.named_parameters():
            if "lora_B" in n:
                p.copy_((torch.randn(p.shape, generator=g) * 0.05).to(p.dtype))
    return model


def rows(lengths, seed: int = 1) -> list[list[int]]:
    import torch

    g = torch.Generator().manual_seed(seed)
    return [[2] + torch.randint(10, VOCAB - 1, (n - 1,), generator=g).tolist() for n in lengths]


def flat_grads(model):
    import torch

    return torch.cat([p.grad.detach().float().flatten() for _, p in model.named_parameters()
                      if p.requires_grad and p.grad is not None])


def forward_backward(model, batch, labels, device, split=None, weights=None):
    import torch

    from .data import step_counts
    from .model import pad, step_loss

    model.zero_grad(set_to_none=True)
    counts = step_counts(labels)
    parts = split or [list(range(len(batch)))]
    for part in parts:
        ids, mask, last = pad([batch[i] for i in part], 0, device)
        y = torch.tensor([labels[i] for i in part], device=device)
        loss, _ = step_loss(model(ids, mask, last), y, counts, weights)
        loss.backward()
    return flat_grads(model)


def rel_err(a, b) -> float:
    return float((a - b).abs().max() / b.abs().max().clamp_min(1e-12))


def cat_logits(model, batch, device):
    import torch

    from .model import pad

    with torch.no_grad():
        return torch.cat([lg.float() for lg in model(*pad(batch, 0, device))], 1)


def run_checks(tmp: Path, device, real_model: str | None = None) -> dict:
    import torch

    from .model import LORA_TARGETS, build, load_backbone, save_checkpoint
    from .taxonomy import Taxonomy

    tax = Taxonomy.default()
    out: dict = {}
    path = tmp / "tiny"
    full = make_tiny(path, real_model)
    # 1. text-only load
    lm, info = load_backbone(path, torch.device("cpu"), dtype=torch.float32)
    ids = torch.tensor([[2, 50, 60, 70, 0, 0]])
    mask = torch.tensor([[1, 1, 1, 1, 0, 0]])
    with torch.no_grad():
        err = float((lm.model(input_ids=ids, attention_mask=mask).last_hidden_state
                     - full.model.language_model(input_ids=ids, attention_mask=mask).last_hidden_state).abs().max())
    ignored = sorted(info["unexpected_keys"])
    out["text_only_load"] = {"missing": sorted(info["missing_keys"]), "ignored_examples": ignored[:4],
                             "ignored": len(ignored), "max_abs_diff_vs_full": err,
                             "pass": not info["missing_keys"] and err < 1e-5}
    del lm, full
    # 2. heads and LoRA
    model = build_tiny(path, device)
    shapes = [tuple(h.weight.shape) for h in model.heads]
    hidden = model.heads[0].in_features
    wrapped = sorted({n.split(".")[-1] for n, m in model.backbone.named_modules() if hasattr(m, "lora_A")})
    out["heads"] = {"shapes": shapes, "lora_modules": wrapped,
                    "pass": shapes == [(k, hidden) for k in tax.sizes()] and set(wrapped) == set(LORA_TARGETS)}
    tol = 1e-4
    # 3. padding invariance, both poolings, in float64: a masking or pooling bug gives errors of order 1, while CUDA
    # fp32 kernels picked per batch shape differ by ~1e-4 on this tiny random model (H100, 2026-10-19: 1.02e-4 in
    # fp32, 2.5e-7 in float64 — the rest is the model's own float32 norm upcasts).
    pad_ok = {}
    for pool in ("last", "mean"):
        m = build_tiny(path, device, pool=pool, dtype="float64")
        m.eval()
        batch = rows([23, 71, 5], seed=2)
        alone = cat_logits(m, [batch[0]], device)
        together = cat_logits(m, batch, device)[0:1]
        pad_ok[pool] = rel_err(alone, together)
    out["padding_invariance"] = {"rel_err": pad_ok, "pass": all(v < tol for v in pad_ok.values())}
    # 4. loss split
    model.train()
    batch = rows([30, 31, 60, 12], seed=3)
    labels = [[0, 1, 2, 0, 1, 2, 3, 1, 0, 2], [1, -1, -1, 1, 0, 0, 11, 3, 5, 0],
              [2, 4, 4, -1, 2, 1, 0, 0, 1, 1], [0, 0, 0, 3, 3, 3, 5, 2, 2, 3]]
    torch.manual_seed(0)
    g_one = forward_backward(model, batch, labels, device, weights=tax.weights())
    torch.manual_seed(0)
    g_two = forward_backward(model, batch, labels, device, split=[[0, 1], [2, 3]], weights=tax.weights())
    split_err = rel_err(g_two, g_one)
    cos = float(torch.nn.functional.cosine_similarity(g_one, g_two, dim=0))
    # fp32 with different padded widths per micro-batch: tiny summation-order differences, not a maths error
    # (a wrong per-step normalisation gives errors of order 1).
    out["loss_split"] = {"grad_rel_err": split_err, "grad_cosine": cos, "pass": split_err < 1e-3 and cos >= 0.9999}
    # 5. save/load
    model.eval()
    before = cat_logits(model, batch, device)
    save_checkpoint(model, tmp / "ckpt", tax.names, {"test": True})
    again = build(tiny_cfg(path), tax.sizes(), device, train=False, adapter_dir=tmp / "ckpt" / "adapter",
                  heads_path=tmp / "ckpt" / "heads.safetensors")
    after = cat_logits(again, batch, device)
    out["save_load"] = {"rel_err": rel_err(after, before), "pass": rel_err(after, before) < tol}
    # 6. CPU/GPU parity
    if torch.device(device).type == "cuda":
        cpu = build(tiny_cfg(path), tax.sizes(), torch.device("cpu"), train=False,
                    adapter_dir=tmp / "ckpt" / "adapter", heads_path=tmp / "ckpt" / "heads.safetensors")
        e = rel_err(cat_logits(cpu, batch, torch.device("cpu")), after.cpu())
        out["cpu_gpu_parity"] = {"rel_err": e, "pass": e < 1e-3}
    else:
        out["cpu_gpu_parity"] = "skipped (no CUDA device)"
    out["ok"] = all(v["pass"] for v in out.values() if isinstance(v, dict))
    return out


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("--model", help="Real checkpoint directory (only its config is used)")
    p.add_argument("--device", default="cuda")
    p.add_argument("--out")
    args = p.parse_args(argv)
    import torch

    device = torch.device(args.device if args.device != "cuda" or torch.cuda.is_available() else "cpu")
    with tempfile.TemporaryDirectory() as tmp:
        out = run_checks(Path(tmp), device, args.model)
    out["device"] = str(device)
    if args.device == "cuda" and device.type != "cuda":
        out["ok"] = False
        out["error"] = "CUDA requested but not available"
    text = json.dumps(out, indent=1, default=str)
    print(text)
    if args.out:
        Path(args.out).write_text(text + "\n")
    return 0 if out["ok"] else 1


if __name__ == "__main__":
    raise SystemExit(main())