"""Tiny-model tests on the pod, before the 9B run. Run as ``python -m classifier.selftest --model DIR``.

The tiny model is the real Qwen3.5 text config cut to four layers (full, DeltaNet, DeltaNet, full attention)
with a small vocabulary and FFN, but the **real head counts and head sizes**, which the kernels constrain.
It is randomly initialised; nothing here needs the 9B weights. Checks (exit 1 if any fails):
1. gradient parity: logits and every LoRA gradient (heads excluded) with the fla + causal-conv1d kernels against
   the PyTorch path, on a right-padded batch: worst per-tensor cosine and norm ratio (fp32 if the kernels
   accept it, otherwise bf16 with looser bounds);
2. padding invariance: a row alone and the same row right-padded in a batch give the same head logits;
3. loss split: one step as one micro-batch and as two micro-batches give the same gradients;
4. save/load: an adapter + heads checkpoint reloads to the same logits.
"""
from __future__ import annotations

import argparse
import copy
import importlib
import json
import tempfile
from pathlib import Path

VOCAB = 2048
LORA = dict(lora_r=8, lora_alpha=8, lora_dropout=0.0, head_dropout=0.0)


def tiny_config(model_dir):
    from .model import text_config
    cfg = copy.deepcopy(text_config(model_dir))
    # full -> linear -> linear -> full: LoRA weights sit upstream of each Gated DeltaNet layer, so a wrong
    # kernel backward (including through g/beta) shows up in trainable gradients.
    cfg.num_hidden_layers = 4
    if hasattr(cfg, 'layer_types'):
        cfg.layer_types = ['full_attention', 'linear_attention', 'linear_attention', 'full_attention']
    cfg.vocab_size = VOCAB
    cfg.hidden_size = 1024
    cfg.intermediate_size = 1024
    if hasattr(cfg, 'mtp_num_hidden_layers'):
        cfg.mtp_num_hidden_layers = 0
    cfg.tie_word_embeddings = False
    cfg.eos_token_id, cfg.pad_token_id = 1, 0
    return cfg


def make_tiny(model_dir, path: Path, dtype):
    import torch
    from transformers import Qwen3_5ForCausalLM
    torch.manual_seed(0)
    lm = Qwen3_5ForCausalLM(tiny_config(model_dir)).to(dtype)
    lm.save_pretrained(str(path))
    return path


def batch(lengths, device, seed=1):
    import torch
    g = torch.Generator().manual_seed(seed)
    return [torch.randint(10, VOCAB, (n,), generator=g).tolist() for n in lengths]


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


def lora_grads(model) -> dict:
    """LoRA gradients by parameter name (heads excluded: they sit after the backbone and would dominate)."""
    return {n: p.grad.detach().float().flatten().clone() for n, p in model.named_parameters()
            if p.requires_grad and p.grad is not None and 'lora_' in n}


def per_tensor(fast: dict, ref: dict) -> dict:
    """Worst per-tensor cosine and norm ratio over LoRA tensors with a non-negligible reference gradient."""
    import torch
    top = max(float(v.norm()) for v in ref.values())
    worst_cos, worst_ratio, n = 1.0, 0.0, 0
    for name, r in ref.items():
        if float(r.norm()) < 1e-6 * top:
            continue
        f = fast[name]
        worst_cos = min(worst_cos, float(torch.nn.functional.cosine_similarity(f, r, dim=0)))
        worst_ratio = max(worst_ratio, abs(float(f.norm()) / float(r.norm()) - 1.0))
        n += 1
    return {'tensors': n, 'worst_cosine': worst_cos, 'worst_norm_ratio_error': worst_ratio}


def cosine(a, b) -> float:
    import torch
    return float(torch.nn.functional.cosine_similarity(a, b, dim=0))


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


def build_tiny(path, device, dtype_name):
    """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
    cfg = dict(LORA, model=str(path), dtype=dtype_name)
    model = build(cfg, [3, 5, 11], device, train=True)
    g = torch.Generator(device='cpu').manual_seed(7)
    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.02).to(p.dtype))
    return model


def forward_backward(model, rows, labels, device, pad_id=0, split=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(rows)))]
    logits_out = []
    for part in parts:
        ids, mask, last = pad([rows[i] for i in part], pad_id, device)
        y = torch.tensor([labels[i] for i in part], device=device)
        logits = model(ids, mask, last)
        loss, _ = step_loss(logits, y, counts)
        loss.backward()
        logits_out.append(torch.cat([lg.detach().float() for lg in logits], dim=1))
    return torch.cat(logits_out, dim=0), flat_grads(model)


def parity(path, device, dtype_name) -> dict:
    import torch
    from .check_env import fast_path, torch_path
    model = build_tiny(path, device, dtype_name)
    model.train()
    fp = fast_path(model)
    rows = batch([37, 64, 9, 120], device)
    labels = [[0, 1, 2], [1, -1, 10], [2, 4, 0], [-1, 0, 3]]
    torch.manual_seed(0)
    logits_fast, _ = forward_backward(model, rows, labels, device)
    lora_fast = lora_grads(model)
    with torch_path():
        torch.manual_seed(0)
        logits_ref, _ = forward_backward(model, rows, labels, device)
        lora_ref = lora_grads(model)
    flat_fast = torch.cat(list(lora_fast.values()))
    flat_ref = torch.cat([lora_ref[k] for k in lora_fast])
    return dict({'dtype': dtype_name, 'fast_path_before': fp, 'logits_rel_err': rel_err(logits_fast, logits_ref),
                 'lora_grad_cosine': cosine(flat_fast, flat_ref), 'lora_grad_rel_err': rel_err(flat_fast, flat_ref)},
                **per_tensor(lora_fast, lora_ref))


def main(argv=None):
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument('--model', required=True, help='Real checkpoint directory (only its config is used)')
    p.add_argument('--device', default='cuda')
    p.add_argument('--skip-parity', action='store_true', help='Local CPU dry run only: no kernels to compare')
    p.add_argument('--out')
    args = p.parse_args(argv)
    import torch
    from .model import build, save_checkpoint, pad
    device = torch.device(args.device)
    out, ok = {}, True
    with tempfile.TemporaryDirectory() as tmp:
        tmp = Path(tmp)
        dtype_name = 'float32'
        if args.skip_parity:
            out['parity'] = 'skipped (--skip-parity)'
            ok = False if device.type == 'cuda' else ok
        else:
            # 1. gradient parity
            try:
                r = parity(make_tiny(args.model, tmp / 'tiny32', torch.float32), device, 'float32')
                r['pass'] = (r['logits_rel_err'] < 1e-3 and r['lora_grad_rel_err'] < 1e-2 and r['worst_cosine'] >= 0.999
                             and r['worst_norm_ratio_error'] <= 0.01)
            except Exception as e:  # kernels may refuse fp32: retry in bf16 with looser bounds
                out['parity_float32_error'] = ('%s: %s' % (type(e).__name__, e))[:500]
                r = parity(make_tiny(args.model, tmp / 'tiny16', torch.bfloat16), device, 'bfloat16')
                r['pass'] = (r['logits_rel_err'] < 5e-2 and r['worst_cosine'] >= 0.99
                             and r['worst_norm_ratio_error'] <= 0.05)
            r['fast_path_was_active'] = bool(r['fast_path_before']['ok'])
            r['pass'] = r['pass'] and r['fast_path_was_active']
            out['parity'] = r
            ok &= r['pass']
            dtype_name = r['dtype']
        # 2-4 on the default path, in fp32 (bf16 if the kernels refused fp32 above)
        path = tmp / ('tiny32' if dtype_name == 'float32' else 'tiny16')
        if not path.exists():
            make_tiny(args.model, path, torch.float32 if dtype_name == 'float32' else torch.bfloat16)
        tol = 1e-4 if dtype_name == 'float32' else 3e-2
        model = build_tiny(path, device, dtype_name)
        model.eval()
        rows = batch([23, 71, 5], device, seed=2)
        with torch.no_grad():
            alone = torch.cat([lg.float() for lg in model(*pad([rows[0]], 0, device))], 1)
            together = torch.cat([lg.float() for lg in model(*pad(rows, 0, device))], 1)[0:1]
        err = rel_err(alone, together)
        out['padding_invariance'] = {'rel_err': err, 'pass': err < tol}
        ok &= out['padding_invariance']['pass']
        model.train()
        rows = batch([30, 31, 60, 12], device, seed=3)
        labels = [[0, 1, 2], [1, -1, 10], [2, 4, -1], [0, 0, 3]]
        torch.manual_seed(0)
        _, g_one = forward_backward(model, rows, labels, device)
        torch.manual_seed(0)
        _, g_two = forward_backward(model, rows, labels, device, split=[[0, 1], [2, 3]])
        out['loss_split'] = {'grad_cosine': cosine(g_one, g_two), 'grad_rel_err': rel_err(g_two, g_one),
                             'pass': (cosine(g_one, g_two) >= (0.9999 if dtype_name == 'float32' else 0.999)
                                      and rel_err(g_two, g_one) < tol * 10)}
        ok &= out['loss_split']['pass']
        model.eval()
        with torch.no_grad():
            before = torch.cat([lg.float() for lg in model(*pad(rows, 0, device))], 1)
        save_checkpoint(model, tmp / 'ckpt', {'test': True})
        cfg = dict(LORA, model=str(path), dtype=dtype_name)
        again = build(cfg, [3, 5, 11], device, train=False, adapter_dir=tmp / 'ckpt' / 'adapter',
                      heads_path=tmp / 'ckpt' / 'heads.pt')
        again.eval()
        with torch.no_grad():
            after = torch.cat([lg.float() for lg in again(*pad(rows, 0, device))], 1)
        err = rel_err(after, before)
        out['save_load'] = {'rel_err': err, 'pass': err < tol}
        ok &= out['save_load']['pass']
    out['ok'] = bool(ok)
    text = json.dumps(out, indent=1, default=str)
    print(text)
    if args.out:
        Path(args.out).write_text(text + '\n')
    return 0 if ok else 1


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