"""Compare generated text across the degradation experiments.

Samples the same prompt(s) from every config that has a trained checkpoint,
using the same RNG seed for each so differences reflect the model, not sampling
luck. Prints the results side by side and writes reports/compare.md.

Usage:
    python scripts/compare.py                              # default prompt, all configs
    python scripts/compare.py --prompt "The dog and the cat"
    python scripts/compare.py --suite                      # every suite prompt
    python scripts/compare.py --configs v1,small,ctx       # a subset
"""

import argparse
import sys
from pathlib import Path

import torch

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "src"))

from sample import (  # noqa: E402
    SUITE_PROMPTS,
    TOK_PATH,
    load_model,
    pick_device,
    sample_text,
)
from tokenizer import BPETokenizer  # noqa: E402

# Baseline first, then the four one-parameter degradations.
DEFAULT_CONFIGS = ["v1", "small", "short", "ctx", "data"]

LABELS = {
    "v1": "v1 (baseline)",
    "small": "small (less capacity)",
    "short": "short (less training)",
    "ctx": "ctx (smaller context)",
    "data": "data (less data)",
}


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--configs", default=",".join(DEFAULT_CONFIGS))
    ap.add_argument("--prompt", default="Once upon a time,")
    ap.add_argument("--suite", action="store_true")
    ap.add_argument("--max-new-tokens", type=int, default=200)
    ap.add_argument("--temperature", type=float, default=0.8)
    ap.add_argument("--top-k", type=int, default=200)
    ap.add_argument("--seed", type=int, default=1234)
    args = ap.parse_args()

    configs = [c.strip() for c in args.configs.split(",") if c.strip()]
    prompts = SUITE_PROMPTS if args.suite else [args.prompt]

    device = pick_device()
    tok = BPETokenizer.load(TOK_PATH)

    # Load each available model once; warn about any not yet trained.
    models = {}
    for name in configs:
        if (ROOT / "checkpoints" / name / "ckpt.pt").exists():
            models[name] = load_model(name, device)
        else:
            print(f"[skip {name}: no checkpoint — train it first]")
    if not models:
        sys.exit("no trained checkpoints found among: " + ", ".join(configs))

    out = ["# Degradation comparison", ""]
    for prompt in prompts:
        print(f"\n########## prompt: {prompt!r} ##########")
        out += [f"## Prompt: {prompt}", ""]
        for name, model in models.items():
            torch.manual_seed(args.seed)  # same RNG state per model = fair compare
            text = sample_text(model, tok, prompt, device,
                               args.max_new_tokens, args.temperature, args.top_k)
            label = LABELS.get(name, name)
            print(f"\n----- {label} -----\n{text}")
            out += [f"### {label}", "", f"> {text}", ""]

    out_path = ROOT / "reports" / "compare.md"
    out_path.parent.mkdir(exist_ok=True)
    out_path.write_text("\n".join(out))
    print(f"\nsaved {out_path}")


if __name__ == "__main__":
    main()