"""Acceptance test for dory/fast_generate.py (KV-cached decoding).

For each loop count, decode N tokens with the cached path (recording the logits used
to pick every token), then run the port's cache-free forward over prompt + generated
tokens and compare the next-token logits at every decoded position.

  .venv/bin/python -I scripts/test_fast_generate.py --tiny                 # random small model, window 8
  .venv/bin/python -I scripts/test_fast_generate.py --model /path/to/Dory-Beam-4B --tokens 300

Both run on CPU in float32 by default.
"""

import argparse
import sys
import time
from pathlib import Path

import torch

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
from dory.fast_generate import generate_stream  # noqa: E402
from dory.modeling import DoryConfig, DoryForCausalLM  # noqa: E402


def tiny_model(seed=0):
    """Small random nGPT-shaped Dory: 2 input, 4 recurrent, 2 output layers, window 8."""
    torch.manual_seed(seed)
    cfg = DoryConfig(vocab_size=257, hidden_size=64, num_hidden_layers=8, intermediate_size=128,
                     num_attention_heads=4, num_key_value_heads=2, head_dim=16, sliding_window=8,
                     rope_parameters={"sliding_attention": {"rope_theta": 10000},
                                      "full_attention": {"rope_theta": 10000000}},
                     layer_types=["sliding_attention", "mlp", "full_attention", "mlp",
                                  "sliding_attention", "mlp", "sliding_attention", "mlp"],
                     n_input_layers=2, n_recurrent_layers=4, n_output_layers=2)
    model = DoryForCausalLM(cfg)
    for p in model.parameters():
        torch.nn.init.normal_(p, std=0.5)
    model.normalize_matrices_()
    # nGPT scalar gains at realistic effective scales (a random init makes softmaxes one-hot
    # and the comparison chaotic): q scale sqrt(d), unit gate/suv, slerp lr 0.05, logit scale 10.
    with torch.no_grad():
        for layer in model.backbone.layers:
            mx = layer.mixer
            if layer.is_attention:
                mx.sqk.fill_(cfg.head_dim**0.25 / mx.sqk_factor)
                mx.s_gate.fill_(1.0 / mx.s_gate_factor)
                layer.attn_alpha.fill_(0.05 / layer.alpha_scale)
            else:
                mx.suv_gate.fill_(1.0 / mx.suv_factor)
                mx.suv_up.fill_(1.0 / mx.suv_factor)
                layer.mlp_alpha.fill_(0.05 / layer.alpha_scale)
        model.lm_head.logit_scale.fill_(10 * cfg.initializer_range)
    return model.float().eval()


def run(model, prompt_ids, loops, n_tokens, tol, seed):
    t0 = time.time()
    stream = generate_stream(model, prompt_ids, n_tokens, n_loops=loops, temperature=1.0, top_p=0.95,
                             eos_ids=(), seed=seed, return_logits=True)
    toks, cached = [], []
    for tok, lg in stream:
        toks.append(tok)
        cached.append(lg[0])
    t_cached = time.time() - t0
    cached = torch.stack(cached)  # [N, V]; row i = logits that produced generated token i
    full_ids = torch.cat([prompt_ids, torch.tensor([toks], dtype=torch.long)], dim=1)
    t0 = time.time()
    ref = model(full_ids, n_loops=loops)[0]  # [T+N, V]
    t_full = time.time() - t0
    p = prompt_ids.shape[1]
    ref = ref[p - 1 : p - 1 + len(toks)]
    diff = (cached - ref).abs()
    scale = ref.abs().max().item()
    print(f"loops={loops}: {len(toks)} tokens, seq {full_ids.shape[1]}, max|Δlogit|={diff.max():.3e} "
          f"(mean {diff.mean():.2e}, logit scale {scale:.1f}), argmax agree "
          f"{(cached.argmax(-1) == ref.argmax(-1)).float().mean():.4f}; "
          f"cached {t_cached:.1f}s ({len(toks) / t_cached:.2f} tok/s), full forward {t_full:.1f}s")
    assert diff.max() < tol, f"max abs diff {diff.max():.3e} exceeds tolerance {tol}"
    return diff.max().item()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", help="directory with config.json + safetensors")
    ap.add_argument("--tiny", action="store_true", help="random small model with sliding window 8")
    ap.add_argument("--tokens", type=int, default=300)
    ap.add_argument("--loops", default="1,2,4")
    ap.add_argument("--tol", type=float, default=None)
    ap.add_argument("--prompt", default="What is 17*23? Think step by step.")
    ap.add_argument("--threads", type=int, default=0)
    ap.add_argument("--float64", action="store_true", help="tiny model in float64 (isolates logic bugs)")
    args = ap.parse_args()
    if args.threads:
        torch.set_num_threads(args.threads)
    torch.manual_seed(0)
    if args.tiny:
        model = tiny_model()
        if args.float64:
            model = model.double()
        prompt_ids = torch.randint(0, 257, (1, 20))  # longer than the window: exercises prefill eviction
        tol = args.tol or (1e-10 if args.float64 else 1e-4)
    else:
        from transformers import PreTrainedTokenizerFast
        model = DoryForCausalLM.from_pretrained(args.model, dtype=torch.float32).eval()
        tok = PreTrainedTokenizerFast.from_pretrained(args.model)
        text = tok.apply_chat_template([{"role": "user", "content": args.prompt}], tokenize=False,
                                       add_generation_prompt=True)
        prompt_ids = tok(text, add_special_tokens=False, return_tensors="pt")["input_ids"]
        tol = args.tol or 5e-3  # measured 1.1e-3 on the released weights (logit scale ~60)
    print(f"prompt tokens: {prompt_ids.shape[1]}, decode tokens: {args.tokens}, tolerance {tol}")
    worst = 0.0
    for loops in (int(x) for x in args.loops.split(",")):
        worst = max(worst, run(model, prompt_ids, loops, args.tokens, tol, seed=loops))
    print(f"PASS: worst max abs logit diff {worst:.3e} < {tol}")


if __name__ == "__main__":
    main()