"""Generate from the model before and after preference tuning, and check for regressions.

Purpose: the measurement half of the DPO lab. Runs the same held-out prompts through the
    starting model and the preference-tuned one at temperature zero, writes both sets of
    answers in the results shape Part 10's judge.py reads, and separately replays your
    Part 10 task set through both models with the deterministic checks, so a style win is
    never reported without asking what it cost everywhere else.
Platform: all. Both models are loaded locally through transformers and PEFT, so no server
    is needed and Track M runs the same path as everyone else.
Minimum memory: 16 GB for a 1B to 4B model held twice in turn; the script loads one model
    at a time and releases it before loading the next.
Assumes: torch, transformers and peft installed; a DPO adapter directory from train-dpo.py;
    optionally the tasks file from Part 10's evaluation lab for the regression check;
    runlog.py sits next to this file.

Usage: python3 compare-before-after.py --model Qwen/Qwen3-1.7B --after runs/dpo-style \
           --out-prefix results --labbook labbook.md
       python3 compare-before-after.py --model Qwen/Qwen3-1.7B --before runs/sft-my-format \
           --after runs/dpo-style --tasks tasks-template.json --labbook labbook.md
"""

from __future__ import annotations

import argparse
import gc
import json
import time
from pathlib import Path
from typing import Any, Optional

import torch

import runlog

# Prompts held out of the preference set on purpose. If the model only improves on the
# prompts it was tuned on, nothing useful happened.
HELD_OUT_PROMPTS = [
    "Explain why two runs of the same model with the same prompt can differ.",
    "What is the practical difference between a base model and an instruct model?",
    "Someone asks whether more context always helps. What do you say?",
    "Explain what an adapter is to someone who has downloaded one and does not know what to do with it.",
    "Describe the trade-off between batch size and latency on a single machine.",
    "How would you check that a converted model still behaves like the one you trained?",
    "Explain in plain words why a model can be fluent and wrong at the same time.",
    "What belongs in a model card, and what usually is not there but should be?",
    "A server answers quickly at first and then slows down. Where would you look?",
    "Explain why the same model gives different answers under two different chat templates.",
    "What is the cheapest useful evaluation you can run on a model you just downloaded?",
    "Describe when running a model on the CPU is a reasonable choice.",
]


def load_model(model_id: str, adapter: Optional[str]):
    """Load the base model, optionally with an adapter, on the best available device."""
    from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: PLC0415

    device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
    dtype = torch.bfloat16 if device == 'cuda' and torch.cuda.is_bf16_supported() else torch.float32
    print(f'loading on {device} with {dtype}')
    if adapter:
        from peft import AutoPeftModelForCausalLM  # noqa: PLC0415
        model = AutoPeftModelForCausalLM.from_pretrained(adapter, dtype=dtype).to(device)
        tokenizer = AutoTokenizer.from_pretrained(adapter)
    else:
        model = AutoModelForCausalLM.from_pretrained(model_id, dtype=dtype).to(device)
        tokenizer = AutoTokenizer.from_pretrained(model_id)
    model.eval()
    return model, tokenizer


def release() -> None:
    """Give the memory back before loading the second model."""
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    elif torch.backends.mps.is_available():
        torch.mps.empty_cache()


def generate(model, tokenizer, prompt: str, max_new_tokens: int) -> tuple[str, float]:
    messages = [{"role": "user", "content": prompt}]
    inputs = tokenizer.apply_chat_template(
        messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt"
    ).to(model.device)
    began = time.time()
    with torch.no_grad():
        out = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
    text = tokenizer.decode(out[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True).strip()
    return text, time.time() - began


def deterministic_checks(task: dict, answer: str) -> dict:
    """Part 10's cheap checks, repeated here so the regression check needs no server."""
    lowered = answer.lower()
    required = [s for s in task.get("must_contain", []) if s.lower() not in lowered]
    forbidden = [s for s in task.get("must_not_contain", []) if s.lower() in lowered]
    words = len(answer.split())
    limit = task.get("max_words")
    return {
        "missing_required": required,
        "found_forbidden": forbidden,
        "words": words,
        "over_length": bool(limit and words > limit),
        "passed": not required and not forbidden and not (limit and words > limit),
    }


def run_one_side(label: str, model_id: str, adapter: Optional[str], style_prompts: list[str],
                 tasks: list[dict], style_rubric: str, max_new_tokens: int) -> dict:
    print(f"\n=== {label}: {adapter or model_id} ===")
    model, tokenizer = load_model(model_id, adapter)
    results = []

    for i, prompt in enumerate(style_prompts, start=1):
        answer, seconds = generate(model, tokenizer, prompt, max_new_tokens)
        results.append({
            "id": f"style-{i:02d}", "category": "style", "prompt": prompt,
            "reference": "", "rubric": style_rubric, "answer": answer,
            "checks": {"words": len(answer.split()), "passed": True},
            "seconds": round(seconds, 2), "prompt_tokens": None, "completion_tokens": None,
        })
        print(f"  style-{i:02d}  {seconds:6.2f}s  {len(answer.split()):4d} words")

    passed = 0
    for task in tasks:
        answer, seconds = generate(model, tokenizer, task["prompt"], max_new_tokens)
        checks = deterministic_checks(task, answer)
        passed += int(checks["passed"])
        results.append({
            "id": task["id"], "category": task.get("category", "regression"),
            "prompt": task["prompt"], "reference": task.get("reference", ""),
            "rubric": task.get("rubric", ""), "answer": answer, "checks": checks,
            "seconds": round(seconds, 2), "prompt_tokens": None, "completion_tokens": None,
        })
        mark = "ok " if checks["passed"] else "BAD"
        print(f"  {mark} {task['id']}  {seconds:6.2f}s")

    del model
    release()
    mean_words = sum(r["checks"]["words"] for r in results) / max(1, len(results))
    return {
        "label": label,
        "model": adapter or model_id,
        "results": results,
        "tasks_passed": passed,
        "tasks_total": len(tasks),
        "mean_words": round(mean_words, 1),
    }


def write_results(path: Path, side: dict, model_id: str, max_new_tokens: int, notes: str) -> None:
    """The shape Part 10's judge.py compare reads: a run block and a list of results."""
    payload = {
        "run": {
            "lab": "part-14/lab-dpo-your-model-to-prefer-your-style",
            "run_id": time.strftime("%Y%m%dT%H%M%S"),
            "task_set": "held-out style prompts plus the Part 10 task set",
            "model": side["model"],
            "quant": "none (transformers, unquantised)",
            "engine": "transformers",
            "base_model": model_id,
            "settings": {"temperature": 0.0, "do_sample": False, "max_tokens": max_new_tokens},
            "date": time.strftime("%Y-%m-%d"),
            "notes": notes,
        },
        "results": side["results"],
    }
    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--model", default="Qwen/Qwen3-1.7B", help="the base model both sides share")
    parser.add_argument("--before", default=None,
                        help="the adapter the DPO run started from, or omit for the bare base model")
    parser.add_argument("--after", required=True, help="the adapter train-dpo.py produced")
    parser.add_argument("--prompts", default=None, help="a file of held-out prompts, one per line")
    parser.add_argument("--tasks", default=None,
                        help="your Part 10 tasks file, for the regression check")
    parser.add_argument("--only-category", default=None, help="use one category of the tasks file only")
    parser.add_argument("--max-tasks", type=int, default=12,
                        help="cap the regression tasks so the lab finishes inside its hour")
    parser.add_argument("--style", default="Direct, short, answer first, plain words, no closing offer of help.",
                        help="the rubric the judge is given for the style prompts")
    parser.add_argument("--max-new-tokens", type=int, default=400)
    parser.add_argument("--out-prefix", default="results")
    parser.add_argument("--labbook", default=None)
    parser.add_argument("--notes", default="")
    args = parser.parse_args()

    if args.prompts:
        style_prompts = [line.strip() for line in Path(args.prompts).read_text(encoding="utf-8").splitlines()
                         if line.strip()]
    else:
        style_prompts = list(HELD_OUT_PROMPTS)

    tasks: list[dict] = []
    if args.tasks:
        spec = json.loads(Path(args.tasks).read_text(encoding="utf-8"))
        tasks = [t for t in spec["tasks"]
                 if args.only_category is None or t.get("category") == args.only_category]
        tasks = tasks[: args.max_tasks]
        print(f"regression check: {len(tasks)} tasks from {args.tasks}")
    else:
        print("no --tasks given, so this run measures style only. A style win with no "
              "regression check is half a result.")

    before = run_one_side("before", args.model, args.before, style_prompts, tasks, args.style,
                          args.max_new_tokens)
    after = run_one_side("after", args.model, args.after, style_prompts, tasks, args.style,
                         args.max_new_tokens)

    before_path = Path(f"{args.out_prefix}-before.json")
    after_path = Path(f"{args.out_prefix}-after.json")
    write_results(before_path, before, args.model, args.max_new_tokens, args.notes)
    write_results(after_path, after, args.model, args.max_new_tokens, args.notes)

    summary: dict[str, Any] = {
        "held_out_prompts": len(style_prompts),
        "regression_tasks": len(tasks),
        "checks_passed_before": before["tasks_passed"],
        "checks_passed_after": after["tasks_passed"],
        "mean_words_before": before["mean_words"],
        "mean_words_after": after["mean_words"],
    }
    print("\n" + json.dumps(summary, indent=2))
    print(f"\nwrote {before_path} and {after_path}")
    print("Now judge them head to head with the position swap, using Part 10's judge:")
    print(f"  python3 judge.py compare --results-a {before_path} --results-b {after_path} \\")
    print("      --judge-model <a different, larger model> --labbook labbook.md")
    if tasks and after["tasks_passed"] < before["tasks_passed"]:
        print("\nThe regression check got worse. That is a result, not a failure of the lab: "
              "record it, then lower the learning rate or raise beta and try again.")

    if args.labbook:
        record = runlog.record(
            labbook=args.labbook,
            lab="part-14/compare-before-after",
            model=args.after,
            dataset={"path": args.tasks, "sha256": runlog.file_sha256(args.tasks) if args.tasks else None,
                     "held_out_prompts": len(style_prompts), "regression_tasks": len(tasks)},
            hyperparameters={"before": args.before or args.model, "after": args.after,
                             "temperature": 0.0, "max_new_tokens": args.max_new_tokens},
            seed=0,
            losses={},
            scores=summary,
            config_path=__file__,
            notes=args.notes or None,
        )
        print(f"recorded run {record['run_id']} in {args.labbook}")


if __name__ == "__main__":
    main()