#!/usr/bin/env python3
"""Compare two evaluation runs and say, in order of importance, what was different about them.

Purpose: the diagnostic instrument for the challenge. Given two lm-evaluation-harness output
    directories that disagree, it reports the score difference per metric and then every setting
    that differs between the runs, ranked by how much that setting is known to move a score. It
    does not decide which run is right; it removes the guessing from the list of candidates.
Platform: all (pure Python, standard library only)
Minimum memory: negligible
Assumes: Python 3.9 or later, and two directories (or two results_*.json files) written by
    lm-evaluation-harness. Runs made by run-suite.sh also carry a run-context.json holding the
    quantisation and engine, which the harness cannot see and which this script compares too.

Usage: python3 diff-eval-runs.py --a run-alpha --b run-beta
       python3 diff-eval-runs.py --a run-alpha --b run-beta --samples --labbook labbook.md
"""

from __future__ import annotations

import argparse
import json
import platform
import time
from pathlib import Path
from typing import Optional

# Ranked by how far each is known to move a benchmark score, worst first. The order is the
# diagnostic procedure: the first difference found is almost always the explanation, and looking
# further down the list before resolving it wastes the afternoon.
RANKED_SETTINGS = [
    ("apply_chat_template", "Chat template applied or not. The largest single lever on an "
                            "instruction-tuned model, and it changes every prompt in the run."),
    ("fewshot_as_multiturn", "Few-shot examples as conversation turns or as one block of text."),
    ("num_fewshot", "Number of in-context examples. Overrides the task file's own default."),
    ("gen_kwargs", "Generation settings: temperature, top-p, whether sampling happens at all."),
    ("model", "Which backend was used to reach the model."),
    ("model_args", "Which model, at which endpoint, with which tokeniser."),
    ("limit", "How many items were scored. A limited run is a different measurement."),
    ("batch_size", "Batching changes summation order and can flip close decisions."),
    ("random_seed", "Seeds. Only matters when something is sampled."),
    ("numpy_seed", "Seeds. Only matters when something is sampled."),
    ("torch_seed", "Seeds. Only matters when something is sampled."),
    ("fewshot_seed", "Which few-shot examples were drawn."),
    ("device", "Where the model ran, when the harness held the weights itself."),
]

CONTEXT_KEYS = [
    ("quant", "The quantisation actually served. The endpoint cannot report it, so it is "
              "recorded by hand and is therefore the field most often wrong."),
    ("engine", "The engine behind the endpoint."),
    ("engine_version", "The engine's version. Sampler and template handling change between them."),
    ("endpoint", "Which server answered."),
    ("alias", "Which gateway alias was called, which decides which file was loaded."),
]


def newest_results(path: Path) -> Optional[Path]:
    if path.is_file():
        return path
    candidates = sorted(path.rglob("results_*.json"))
    return candidates[-1] if candidates else None


def load_run(path: Path) -> dict:
    results_file = newest_results(path)
    if results_file is None:
        raise SystemExit(f"no results_*.json found under {path}")
    payload = json.loads(results_file.read_text(encoding="utf-8"))

    context = {}
    for parent in [path] + list(path.parents)[:2]:
        candidate = parent / "run-context.json"
        if candidate.is_file():
            context = json.loads(candidate.read_text(encoding="utf-8"))
            break

    return {"file": results_file, "payload": payload, "context": context}


def scores(payload: dict) -> dict:
    out = {}
    for task, metrics in (payload.get("results") or {}).items():
        if not isinstance(metrics, dict):
            continue
        for key, value in metrics.items():
            if isinstance(value, (int, float)) and not key.startswith("alias") \
                    and "_stderr" not in key:
                out[f"{task}/{key}"] = float(value)
    return out


def show(value) -> str:
    text = json.dumps(value) if not isinstance(value, str) else value
    return text if len(text) <= 110 else text[:107] + "..."


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--a", required=True, help="first run: a directory or a results_*.json")
    parser.add_argument("--b", required=True, help="second run: a directory or a results_*.json")
    parser.add_argument("--samples", action="store_true",
                        help="also compare the first logged sample from each run, which shows the "
                             "prompt as actually sent")
    parser.add_argument("--out", default=None)
    parser.add_argument("--labbook", default=None)
    args = parser.parse_args()

    run_a = load_run(Path(args.a))
    run_b = load_run(Path(args.b))
    print(f"A: {run_a['file']}")
    print(f"B: {run_b['file']}")

    # ---- 1. what actually differs in the numbers -------------------------
    score_a, score_b = scores(run_a["payload"]), scores(run_b["payload"])
    shared = sorted(set(score_a) & set(score_b))
    print("\n== Scores ==")
    if not shared:
        print("  The two runs share no metric. They did not measure the same thing, which is "
              "already the answer.")
    differences = []
    for key in shared:
        delta = score_b[key] - score_a[key]
        differences.append({"metric": key, "a": score_a[key], "b": score_b[key],
                            "delta": round(delta, 4)})
        flag = "  <-- " if abs(delta) >= 0.01 else "      "
        print(f"  {key:52s} A={score_a[key]:.4f}  B={score_b[key]:.4f}  Δ={delta:+.4f}{flag}")
    only_a = sorted(set(score_a) - set(score_b))
    only_b = sorted(set(score_b) - set(score_a))
    for key in only_a:
        print(f"  {key:52s} present in A only")
    for key in only_b:
        print(f"  {key:52s} present in B only")

    # ---- 2. the settings, in the order they matter -----------------------
    config_a = run_a["payload"].get("config") or {}
    config_b = run_b["payload"].get("config") or {}
    found = []
    print("\n== Settings that differ, most consequential first ==")
    for key, why in RANKED_SETTINGS:
        left, right = config_a.get(key), config_b.get(key)
        if left != right:
            found.append({"setting": key, "a": left, "b": right, "why": why})
            print(f"  {key}")
            print(f"      A: {show(left)}")
            print(f"      B: {show(right)}")
            print(f"      {why}")

    other_keys = (set(config_a) | set(config_b)) - {k for k, _ in RANKED_SETTINGS}
    for key in sorted(other_keys):
        if config_a.get(key) != config_b.get(key):
            found.append({"setting": key, "a": config_a.get(key), "b": config_b.get(key),
                          "why": "not on the ranked list; check it after the ones above"})
            print(f"  {key} (unranked)")
            print(f"      A: {show(config_a.get(key))}")
            print(f"      B: {show(config_b.get(key))}")

    # ---- 3. the task configuration itself --------------------------------
    task_a = run_a["payload"].get("configs") or {}
    task_b = run_b["payload"].get("configs") or {}
    print("\n== Task configuration ==")
    task_diffs = []
    for task in sorted(set(task_a) | set(task_b)):
        left, right = task_a.get(task, {}), task_b.get(task, {})
        for field in sorted(set(left) | set(right)):
            if left.get(field) != right.get(field):
                task_diffs.append({"task": task, "field": field,
                                   "a": left.get(field), "b": right.get(field)})
                print(f"  {task}.{field}")
                print(f"      A: {show(left.get(field))}")
                print(f"      B: {show(right.get(field))}")
    if not task_diffs:
        print("  Identical. Both runs used the same task definition, so the difference is in the "
              "run settings or outside the harness entirely.")

    # ---- 4. what the harness could not see -------------------------------
    print("\n== Outside the harness ==")
    context_diffs = []
    if not run_a["context"] and not run_b["context"]:
        print("  Neither run has a run-context.json, so the quantisation, the engine and its "
              "version are unrecorded for both. That is itself a finding: the two runs may have "
              "been served different files and nothing here would show it.")
    else:
        for key, why in CONTEXT_KEYS:
            left, right = run_a["context"].get(key), run_b["context"].get(key)
            if left != right:
                context_diffs.append({"setting": key, "a": left, "b": right, "why": why})
                print(f"  {key}")
                print(f"      A: {show(left)}")
                print(f"      B: {show(right)}")
                print(f"      {why}")
        if not context_diffs:
            print("  Identical where recorded.")

    # ---- 5. the prompt as actually sent ----------------------------------
    if args.samples:
        print("\n== First logged sample from each run ==")
        for name, run in (("A", run_a), ("B", run_b)):
            base = run["file"].parent
            sample_files = sorted(base.glob("samples_*.jsonl"))
            if not sample_files:
                print(f"  {name}: no samples_*.jsonl. Re-run with --log_samples; without it the "
                      f"prompt as sent is not recoverable.")
                continue
            with sample_files[0].open(encoding="utf-8") as handle:
                first = handle.readline()
            try:
                record = json.loads(first)
            except ValueError:
                print(f"  {name}: could not parse the first sample line.")
                continue
            shown = False
            for field in ("arguments", "doc", "filtered_resps", "resps"):
                if field in record:
                    print(f"  {name}.{field}: {show(record[field])}")
                    shown = True
            if not shown:
                print(f"  {name}: fields present: {', '.join(sorted(record)[:12])}")
        print("\n  Read the two prompts side by side. Role markers present in one and absent in "
              "the other is a chat-template difference; a different number of worked examples is "
              "a shot-count difference. Both are visible here and invisible in the score alone.")

    # ---- verdict ---------------------------------------------------------
    print("\n== Verdict ==")
    if found:
        first = found[0]
        print(f"  Highest-ranked difference: {first['setting']}.")
        print("  Change that one thing, re-run, and see whether the gap closes. One change at a "
              "time, or you will not know which one worked.")
    elif context_diffs:
        print(f"  The harness settings are identical; the difference is outside it: "
              f"{context_diffs[0]['setting']}.")
    elif task_diffs:
        print("  The run settings are identical and the task definition is not. Check the harness "
              "version and the dataset revision.")
    else:
        print("  No recorded setting differs. Either the difference is genuine run-to-run "
              "variation - establish the noise floor by running one configuration three times - "
              "or something changed that neither the harness nor run-context.json records, such "
              "as which file the endpoint had loaded.")

    record = {
        "lab": "part-16/challenge-the-benchmark-that-lied",
        "run_id": time.strftime("%Y%m%dT%H%M%S"),
        "a": str(run_a["file"]),
        "b": str(run_b["file"]),
        "score_differences": differences,
        "setting_differences": found,
        "task_differences": task_diffs,
        "context_differences": context_diffs,
        "host": platform.platform(),
        "date": time.strftime("%Y-%m-%d"),
    }
    if args.out:
        Path(args.out).write_text(json.dumps(record, indent=2), encoding="utf-8")
        print(f"\nwritten to {args.out}")
    if args.labbook:
        with Path(args.labbook).open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(record) + "\n")
        print(f"recorded in {args.labbook}")


if __name__ == "__main__":
    main()
