"""Compare a from-scratch model with a fine-tuned pretrained one, on equal terms.

Purpose: the project's evaluation. Two models with different tokenisers and very
         different parameter counts cannot be compared on mean loss, so this scores
         both in bits per byte on exactly the same held-out domain text, and on a
         small task set where the metric is again bits per byte of the expected
         continuation. Both numbers are tokenisation-independent, which is what
         makes the comparison fair rather than merely available.
Platform: all (cuda on Tracks S and N and on Track X with ROCm, mps on Track M, cpu
          anywhere as a slow fallback); the device is autodetected and recorded.
Minimum memory: 8 GB
Assumes: this is run with an interpreter that can import both nanochat (from the
         clone) and transformers and peft; the domain shards exist; a nanochat
         checkpoint exists under $NANOCHAT_BASE_DIR/base_checkpoints/<tag>.

Usage: python compare-models.py --shards ~/.cache/nanochat-domain/base_data_climbmix
                                --nanochat-tag d6-domain
                                [--hf-model Qwen/Qwen3-1.7B-Base]
                                [--adapter domain-lora]
                                [--tasks domain-tasks.json]
                                [--max-val-chars 200000] [--labbook labbook.md]

Bits per byte is summed negative log likelihood in nats, divided by the natural
logarithm of two and by the number of UTF-8 bytes the scored tokens represent. A
model with a larger vocabulary predicts fewer, longer tokens and so wins on mean
loss for free; dividing by bytes removes that advantage entirely.
"""
from __future__ import annotations

import argparse
import json
import math
import os
import platform
import secrets
import sys
from datetime import datetime, timezone
from pathlib import Path

LN2 = math.log(2.0)


# --------------------------------------------------------------------------- #
# Held-out text                                                                #
# --------------------------------------------------------------------------- #

def read_validation_text(shard_dir: Path, max_chars: int) -> str:
    """The validation shard, concatenated, capped so the comparison stays quick."""
    import pyarrow.parquet as pq

    paths = sorted(shard_dir.glob("*.parquet"))
    if len(paths) < 2:
        raise SystemExit(f"Expected at least two parquet shards in {shard_dir}, found {len(paths)}.")
    documents = pq.read_table(paths[-1], columns=["text"]).column("text").to_pylist()
    text = "\n\n".join(documents)
    return text[:max_chars]


# --------------------------------------------------------------------------- #
# Scoring: nanochat                                                            #
# --------------------------------------------------------------------------- #

def nanochat_score_text(model, tokeniser, text: str, window: int) -> tuple[float, int]:
    """Summed nats and scored UTF-8 bytes for a stretch of text."""
    import torch

    ids = tokeniser.encode(text)
    nats, byte_count = 0.0, 0
    for start in range(0, max(0, len(ids) - 1), window):
        chunk = ids[start : start + window + 1]
        if len(chunk) < 2:
            break
        x = torch.tensor([chunk[:-1]], dtype=torch.long, device=model.get_device())
        y = torch.tensor([chunk[1:]], dtype=torch.long, device=model.get_device())
        with torch.no_grad():
            nats += float(model(x, y, loss_reduction="sum").item())
        byte_count += len(tokeniser.decode(chunk[1:]).encode("utf-8"))
    return nats, byte_count


def nanochat_score_continuation(model, tokeniser, prompt: str, expected: str) -> tuple[float, int]:
    """Summed nats and bytes for `expected` given `prompt`, scoring only `expected`."""
    import torch

    prompt_ids = tokeniser.encode(prompt, prepend="<|bos|>")
    expected_ids = tokeniser.encode(expected)
    if not expected_ids:
        return 0.0, 0
    ids = prompt_ids + expected_ids
    x = torch.tensor([ids[:-1]], dtype=torch.long, device=model.get_device())
    y = torch.tensor([ids[1:]], dtype=torch.long, device=model.get_device())
    # Score only the positions that predict the expected continuation.
    mask = torch.full_like(y, -1)
    mask[:, len(prompt_ids) - 1 :] = y[:, len(prompt_ids) - 1 :]
    with torch.no_grad():
        nats = float(model(x, mask, loss_reduction="sum").item())
    return nats, len(expected.encode("utf-8"))


# --------------------------------------------------------------------------- #
# Scoring: a Hugging Face causal model                                         #
# --------------------------------------------------------------------------- #

def hf_score_text(model, tokeniser, text: str, window: int, device) -> tuple[float, int]:
    import torch
    import torch.nn.functional as F

    ids = tokeniser(text, add_special_tokens=False)["input_ids"]
    nats, byte_count = 0.0, 0
    for start in range(0, max(0, len(ids) - 1), window):
        chunk = ids[start : start + window + 1]
        if len(chunk) < 2:
            break
        x = torch.tensor([chunk[:-1]], dtype=torch.long, device=device)
        y = torch.tensor(chunk[1:], dtype=torch.long, device=device)
        with torch.no_grad():
            logits = model(input_ids=x).logits[0].float()
        nats += float(F.cross_entropy(logits, y, reduction="sum").item())
        byte_count += len(tokeniser.decode(chunk[1:]).encode("utf-8"))
    return nats, byte_count


def hf_score_continuation(model, tokeniser, prompt: str, expected: str, device) -> tuple[float, int]:
    import torch
    import torch.nn.functional as F

    prompt_ids = tokeniser(prompt, add_special_tokens=False)["input_ids"]
    expected_ids = tokeniser(expected, add_special_tokens=False)["input_ids"]
    if not expected_ids:
        return 0.0, 0
    ids = prompt_ids + expected_ids
    x = torch.tensor([ids[:-1]], dtype=torch.long, device=device)
    with torch.no_grad():
        logits = model(input_ids=x).logits[0].float()
    target = torch.tensor(expected_ids, dtype=torch.long, device=device)
    scored = logits[len(prompt_ids) - 1 :]
    nats = float(F.cross_entropy(scored, target, reduction="sum").item())
    return nats, len(expected.encode("utf-8"))


# --------------------------------------------------------------------------- #

def bits_per_byte(nats: float, byte_count: int) -> float | None:
    return None if byte_count == 0 else nats / (LN2 * byte_count)


def normalise(text: str) -> str:
    return " ".join(text.lower().split())


def append_run_log(path: str, record: dict) -> None:
    """Append one JSON line in the run-log format Part 11 introduces."""
    notebook = Path(path)
    if not notebook.exists():
        notebook.write_text("# Lab notebook\n\n## Results\n\n", encoding="utf-8")
    with notebook.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(record, ensure_ascii=False) + "\n")


def new_run_id() -> str:
    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    return f"{stamp}-{secrets.token_hex(3)}"


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--shards", required=True, help="directory of domain parquet shards")
    parser.add_argument("--nanochat-tag", required=True, help="checkpoint tag of the from-scratch model")
    parser.add_argument("--nanochat", default=os.environ.get("NANOCHAT", str(Path.home() / "nanochat")))
    parser.add_argument("--nanochat-step", type=int, default=None)
    parser.add_argument("--hf-model", default="Qwen/Qwen3-1.7B-Base")
    parser.add_argument("--adapter", default=None, help="LoRA directory from continue-pretraining.py")
    parser.add_argument("--tasks", default=None, help="JSON list of {prompt, expected} items")
    parser.add_argument("--max-val-chars", type=int, default=200_000)
    parser.add_argument("--window", type=int, default=None, help="scoring window; default: the model's sequence length")
    parser.add_argument("--max-new-tokens", type=int, default=24)
    parser.add_argument("--labbook", default=None)
    parser.add_argument("--device", default="", help="cuda|mps|cpu (empty = autodetect)")
    args = parser.parse_args()

    import torch

    shard_dir = Path(args.shards).expanduser()
    text = read_validation_text(shard_dir, args.max_val_chars)
    print(f"held-out text: {len(text):,} characters, {len(text.encode('utf-8')):,} bytes")

    tasks = []
    if args.tasks:
        tasks = json.loads(Path(args.tasks).read_text(encoding="utf-8"))
        print(f"task set: {len(tasks)} item(s)")

    results: dict[str, dict] = {}

    # ---- the from-scratch model ------------------------------------------
    repo = Path(args.nanochat).expanduser()
    if not (repo / "nanochat").is_dir():
        sys.exit(f"No nanochat package under {repo}. Pass --nanochat or set NANOCHAT.")
    sys.path.insert(0, str(repo))
    from nanochat.checkpoint_manager import load_model
    from nanochat.common import autodetect_device_type, compute_init
    from nanochat.engine import Engine

    device_type = autodetect_device_type() if args.device == "" else args.device
    _, _, _, _, device = compute_init(device_type)
    scratch_model, scratch_tok, scratch_meta = load_model(
        "base", device, phase="eval", model_tag=args.nanochat_tag, step=args.nanochat_step
    )
    window = args.window or int(scratch_meta["model_config"]["sequence_len"])
    print(f"\nscoring the from-scratch model (window {window})")
    nats, byte_count = nanochat_score_text(scratch_model, scratch_tok, text, window)
    scratch = {"val_bpb": bits_per_byte(nats, byte_count), "val_bytes": byte_count, "tasks": []}

    engine = Engine(scratch_model, scratch_tok)
    for item in tasks:
        t_nats, t_bytes = nanochat_score_continuation(
            scratch_model, scratch_tok, item["prompt"], item["expected"]
        )
        prompt_ids = scratch_tok(item["prompt"], prepend="<|bos|>")
        generated, _ = engine.generate_batch(
            prompt_ids, num_samples=1, max_tokens=args.max_new_tokens, temperature=0
        )
        continuation = scratch_tok.decode(generated[0])[len(item["prompt"]) :]
        scratch["tasks"].append({
            "prompt": item["prompt"],
            "expected": item["expected"],
            "bpb": bits_per_byte(t_nats, t_bytes),
            "greedy": continuation,
            "prefix_match": normalise(continuation).startswith(normalise(item["expected"])),
        })
    results["from scratch"] = scratch
    del scratch_model, engine

    # ---- the fine-tuned pretrained model ----------------------------------
    from transformers import AutoModelForCausalLM, AutoTokenizer

    hf_dtype = torch.float32 if device.type == "cpu" else torch.bfloat16
    hf_tok = AutoTokenizer.from_pretrained(args.hf_model)
    try:
        hf_model = AutoModelForCausalLM.from_pretrained(args.hf_model, dtype=hf_dtype)
    except TypeError:
        hf_model = AutoModelForCausalLM.from_pretrained(args.hf_model, torch_dtype=hf_dtype)
    if args.adapter:
        from peft import PeftModel

        hf_model = PeftModel.from_pretrained(hf_model, args.adapter)
    hf_model = hf_model.to(device)
    hf_model.eval()

    label = f"{args.hf_model}{' + adapter' if args.adapter else ' (no adapter)'}"
    print(f"\nscoring {label} (window {window})")
    nats, byte_count = hf_score_text(hf_model, hf_tok, text, window, device)
    reference = {"val_bpb": bits_per_byte(nats, byte_count), "val_bytes": byte_count, "tasks": []}

    for item in tasks:
        t_nats, t_bytes = hf_score_continuation(
            hf_model, hf_tok, item["prompt"], item["expected"], device
        )
        ids = hf_tok(item["prompt"], add_special_tokens=False, return_tensors="pt").to(device)
        with torch.no_grad():
            generated = hf_model.generate(
                **ids, max_new_tokens=args.max_new_tokens, do_sample=False
            )
        continuation = hf_tok.decode(generated[0][ids["input_ids"].shape[1] :], skip_special_tokens=True)
        reference["tasks"].append({
            "prompt": item["prompt"],
            "expected": item["expected"],
            "bpb": bits_per_byte(t_nats, t_bytes),
            "greedy": continuation,
            "prefix_match": normalise(continuation).startswith(normalise(item["expected"])),
        })
    results[label] = reference

    # ---- report -----------------------------------------------------------
    def summarise(entry: dict) -> tuple[float | None, float | None]:
        scores = [t["bpb"] for t in entry["tasks"] if t["bpb"] is not None]
        matches = [t["prefix_match"] for t in entry["tasks"]]
        mean_bpb = sum(scores) / len(scores) if scores else None
        match_rate = sum(matches) / len(matches) if matches else None
        return mean_bpb, match_rate

    print("\n" + "=" * 84)
    print(f"{'Model':<44} {'val bpb':>9} {'task bpb':>9} {'prefix match':>13}")
    print("-" * 84)
    for name, entry in results.items():
        mean_bpb, match_rate = summarise(entry)
        print(
            f"{name[:44]:<44} "
            f"{entry['val_bpb']:>9.4f} "
            f"{(mean_bpb if mean_bpb is not None else float('nan')):>9.4f} "
            f"{(f'{match_rate:.0%}' if match_rate is not None else 'n/a'):>13}"
        )
    print("=" * 84)
    print(
        "Lower bits per byte is better. Both columns are divided by UTF-8 bytes rather\n"
        "than by tokens, so the two tokenisers do not distort the comparison. The\n"
        "parameter counts do: state them next to these numbers, because a comparison\n"
        "between models of different sizes is a statement about what the extra\n"
        "parameters and the extra pretraining bought, not about the method."
    )

    if args.labbook:
        record = {
            "run_id": new_run_id(),
            "lab": "part-12/domain-micro-model/comparison",
            "date": datetime.now(timezone.utc).strftime("%Y-%m-%d"),
            "config_commit": None,
            "model": {"from_scratch": args.nanochat_tag, "reference": args.hf_model,
                      "adapter": args.adapter},
            "dataset": {"name": str(shard_dir), "held_out_characters": len(text),
                        "held_out_bytes": len(text.encode("utf-8"))},
            "hyperparameters": {"window": window, "max_new_tokens": args.max_new_tokens},
            "seed": None,
            "hardware": {"os": f"{platform.system()} {platform.release()}",
                         "arch": platform.machine(), "accelerator": device.type},
            "versions": {"python": platform.python_version(), "torch": torch.__version__},
            "losses": {name: entry["val_bpb"] for name, entry in results.items()},
            "scores": {
                name: {"task_bpb": summarise(entry)[0], "prefix_match": summarise(entry)[1]}
                for name, entry in results.items()
            },
            "notes": {name: entry["tasks"] for name, entry in results.items()},
        }
        append_run_log(args.labbook, record)
        print(f"\nrecorded in {args.labbook}")


if __name__ == "__main__":
    main()