#!/usr/bin/env python3
"""Strict, dependency-free analysis of the pinned Spark MMLU-Pro protocol.

Usage: python3 analyze_mmlu.py --config results/mmlu/analysis-config.json
All input paths are relative to the config file. No partial denominators,
duplicate overwrites, dropped rows, or implicit cross-model comparisons.
"""

import argparse
from collections import Counter, defaultdict
from datetime import datetime, timezone
import hashlib
import json
import math
from pathlib import Path
import sys


class ValidationError(ValueError):
    pass


def require(condition, message):
    if not condition:
        raise ValidationError(message)


def sha256(path):
    digest = hashlib.sha256()
    with Path(path).open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def load_index(index_path, subset_path, expected_n=2000):
    rows = json.loads(Path(index_path).read_text())
    subset = json.loads(Path(subset_path).read_text())
    require(isinstance(rows, list), "Task index must be a JSON list")
    require(len(rows) == expected_n, f"Index has {len(rows)} tasks, expected {expected_n}")
    require(subset.get("n") == expected_n, "Subset count differs from required task count")
    index = {}
    question_ids = []
    for position, row in enumerate(rows):
        require(isinstance(row, dict), f"Index row {position} is not an object")
        for key in ("i", "question_id", "n_options", "answer_index"):
            require(type(row.get(key)) is int, f"Index row {position}: {key} must be an integer")
        i = row["i"]
        require(i not in index, f"Duplicate task index {i}")
        require(i == position, f"Index row {position} has task ID {i}: order must match prompts")
        require(2 <= row["n_options"] <= 10, f"Task {i}: unsupported option count")
        require(0 <= row["answer_index"] < row["n_options"], f"Task {i}: gold label out of range")
        require(isinstance(row.get("category"), str) and row["category"], f"Task {i}: invalid category")
        index[i] = row
        question_ids.append(row["question_id"])
    require(len(set(question_ids)) == expected_n, "Duplicate dataset question ID")
    require(question_ids == subset.get("question_ids"), "Task question IDs/order differ from pinned subset")
    categories = dict(sorted(Counter(r["category"] for r in rows).items()))
    require(categories == subset.get("allocation"), "Category counts differ from pinned subset")
    return index, subset


def load_scores(path, index):
    """Validate every row before returning any scores; accept rounded argmax ties."""
    rows = {}
    ties = []
    path = Path(path)
    with path.open(encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, 1):
            where = f"{path.name}:{line_number}"
            require(bool(line.strip()), f"{where}: blank row")
            fields = line.rstrip("\r\n").split("\t")
            require(len(fields) >= 5, f"{where}: expected task, prediction, gold and option scores")
            try:
                i, pred, gold = map(int, fields[:3])
                logprobs = tuple(map(float, fields[3:]))
            except ValueError as error:
                raise ValidationError(f"{where}: malformed numeric field") from error
            require(i in index, f"{where}: unknown task index {i}")
            require(i not in rows, f"{where}: duplicate task index {i}")
            task = index[i]
            n_options = task["n_options"]
            require(len(logprobs) == n_options, f"{where}: {len(logprobs)} scores, expected {n_options}")
            require(gold == task["answer_index"], f"{where}: gold label {gold} differs from index")
            require(0 <= pred < n_options, f"{where}: prediction out of range")
            require(all(math.isfinite(x) for x in logprobs), f"{where}: nonfinite logprob")
            require(max(logprobs) <= 0, f"{where}: positive logprob")
            require(sum(math.exp(x) for x in logprobs) <= 1.00001, f"{where}: option probabilities exceed one")
            maximum = max(logprobs)
            require(logprobs[pred] == maximum, f"{where}: prediction is not an argmax of recorded logprobs")
            if logprobs.count(maximum) > 1:
                ties.append(i)
            rows[i] = {"prediction": pred, "gold": gold, "logprobs": logprobs}
    missing = sorted(set(index) - set(rows))
    require(not missing, f"{path.name}: missing {len(missing)} tasks; first missing IDs: {missing[:10]}")
    require(len(rows) == len(index), f"{path.name}: invalid total task count")
    return rows, {
        "passed": True,
        "n_rows": len(rows),
        "unique_indices": len(rows),
        "complete_index_coverage": True,
        "gold_labels_match": True,
        "option_counts_match": True,
        "finite_nonpositive_logprobs": True,
        "prediction_is_recorded_argmax": True,
        "rounded_argmax_tie_indices": ties,
        "tie_note": "TSV has rounded logprobs; prediction must attain the recorded maximum. Exact pre-rounding tie-breaking cannot be recovered.",
    }


def summarize(rows, index):
    n = len(rows)
    correct = sum(r["prediction"] == r["gold"] for r in rows.values())
    accuracy = correct / n
    by_category = defaultdict(list)
    for i in rows:
        by_category[index[i]["category"]].append(i)
    per_category = {}
    for category, ids in sorted(by_category.items()):
        c = sum(rows[i]["prediction"] == rows[i]["gold"] for i in ids)
        p = c / len(ids)
        per_category[category] = {
            "n": len(ids), "correct": c, "accuracy_pct": 100 * p,
            "binomial_se_pp": 100 * math.sqrt(p * (1 - p) / len(ids)),
        }
    return {
        "n": n, "correct": correct, "accuracy_pct": 100 * accuracy,
        "binomial_se_pp": 100 * math.sqrt(accuracy * (1 - accuracy) / n),
        "per_category": per_category,
        "prediction_counts": dict(sorted(Counter(str(r["prediction"]) for r in rows.values()).items())),
        "gold_counts": dict(sorted(Counter(str(r["gold"]) for r in rows.values()).items())),
    }


def exact_mcnemar(a_only, b_only):
    require(type(a_only) is int and a_only >= 0, "a_only must be a nonnegative integer")
    require(type(b_only) is int and b_only >= 0, "b_only must be a nonnegative integer")
    n = a_only + b_only
    if not n:
        return 1.0
    # Exact integer arithmetic before final conversion: no factorial overflow.
    numerator = 2 * sum(math.comb(n, k) for k in range(min(a_only, b_only) + 1))
    return min(1.0, numerator / (1 << n))


def compare(a, b):
    require(set(a) == set(b), "Paired comparison requires identical complete task IDs")
    n = len(a)
    both = a_only = b_only = neither = agree = 0
    max_logprob_difference = 0.0
    for i in a:
        require(a[i]["gold"] == b[i]["gold"], f"Paired task {i}: mismatched gold labels")
        require(len(a[i]["logprobs"]) == len(b[i]["logprobs"]), f"Paired task {i}: mismatched options")
        ac = a[i]["prediction"] == a[i]["gold"]
        bc = b[i]["prediction"] == b[i]["gold"]
        both += ac and bc
        a_only += ac and not bc
        b_only += bc and not ac
        neither += not ac and not bc
        agree += a[i]["prediction"] == b[i]["prediction"]
        max_logprob_difference = max(max_logprob_difference, *(abs(x-y) for x, y in zip(a[i]["logprobs"], b[i]["logprobs"])))
    discordant = a_only + b_only
    delta = (a_only - b_only) / n
    return {
        "n": n,
        "delta_definition": "accuracy(a) minus accuracy(b)",
        "delta_pp": 100 * delta,
        "paired_se_pp": 100 * math.sqrt(max(0.0, discordant / n - delta * delta) / (n - 1)) if n > 1 else None,
        "paired_se_definition": "SE of per-question correctness differences using sample variance (n-1)",
        "mcnemar_null_se_pp": 100 * math.sqrt(discordant) / n,
        "both_correct": both, "a_only_correct": a_only, "b_only_correct": b_only, "neither_correct": neither,
        "prediction_agreement_count": agree,
        "prediction_agreement_pct": 100 * agree / n,
        "correctness_agreement_count": both + neither,
        "correctness_agreement_pct": 100 * (both + neither) / n,
        "max_abs_logprob_difference": max_logprob_difference,
        "mcnemar_exact_two_sided_p": exact_mcnemar(a_only, b_only),
        "p_value_adjustment": "none; pairwise descriptive comparisons",
    }


def markdown(report):
    out = ["# MMLU-Pro evaluation results", "", f"Generated {report['generated_at']}. Every included TSV passed strict validation on all {report['dataset']['n']:,} pinned questions.", "",
           "Zero-shot raw prompt; no chat template or reasoning generation. One prompt decode, vocabulary log-softmax at the single-token continuations ` A` through ` J`, argmax over available options.", "",
           f"Uniform random choice accuracy averaged over actual option counts: **{report['dataset']['chance_floor_pct']:.4f}%**. Accuracy SE is binomial; this is a fixed stratified subset, and SE does not include runtime or dataset-selection uncertainty.", "",
           "| Run | Status | Correct | Accuracy | Binomial SE |", "|---|---|---:|---:|---:|"]
    for name, run in report["runs"].items():
        s = run["summary"]
        out.append(f"| {name} | {run['status']} | {s['correct']}/{s['n']} | {s['accuracy_pct']:.2f}% | {s['binomial_se_pp']:.2f} pp |")
    if report["comparisons"]:
        out += ["", "Every delta below is **A minus B**. Prediction agreement compares chosen option indices; correctness agreement is recorded separately in JSON. McNemar p-values are exact, two-sided and unadjusted.", "",
                "| A | B | Context | Delta | Paired SE | A-only / B-only correct | Prediction agreement | McNemar p |", "|---|---|---|---:|---:|---:|---:|---:|"]
        for pair in report["comparisons"]:
            s = pair["statistics"]
            out.append(f"| {pair['a']} | {pair['b']} | {pair['classification']} | {s['delta_pp']:+.2f} pp | {s['paired_se_pp']:.2f} pp | {s['a_only_correct']} / {s['b_only_correct']} | {s['prediction_agreement_pct']:.2f}% | {s['mcnemar_exact_two_sided_p']:.6g} |")
    out += ["", "## Per-category accuracy", "", "Category estimates have small denominators; do not select a runtime using these scores.", "",
            "| Category | n | " + " | ".join(report["runs"]) + " |",
            "|---|---:|" + "---:|" * len(report["runs"])]
    for category, n in report["dataset"]["category_counts"].items():
        values = [f"{run['summary']['per_category'][category]['accuracy_pct']:.2f}%" for run in report["runs"].values()]
        out.append(f"| {category} | {n} | " + " | ".join(values) + " |")
    out += ["", "## Interpretation and provenance", ""]
    out.extend(f"- {note}" for note in report["notes"])
    for pair in report["comparisons"]:
        if pair.get("interpretation"):
            out.append(f"- {pair['a']} versus {pair['b']}: {pair['interpretation']}")
    out += ["", "Input SHA256 hashes, model provenance, validation flags, prediction counts, paired cells and category counts are in the companion JSON. Valid TSV structure does not establish model/runtime equivalence.", ""]
    return "\n".join(out)


def analyze(config_path):
    config_path = Path(config_path).resolve()
    config = json.loads(config_path.read_text())
    base = config_path.parent
    index_path = base / config["index"]
    subset_path = base / config["subset"]
    index, subset = load_index(index_path, subset_path)
    report = {
        "schema_version": 1,
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "analysis_script_sha256": sha256(__file__),
        "config_sha256": sha256(config_path),
        "dataset": {
            "id": subset["dataset"], "split": subset["split"], "n": len(index), "seed": subset["seed"],
            "source_parquet_sha256": subset["file_sha256"],
            "task_index_sha256": sha256(index_path), "subset_sha256": sha256(subset_path),
            "category_counts": dict(sorted(Counter(r["category"] for r in index.values()).items())),
            "option_count_distribution": dict(sorted(Counter(r["n_options"] for r in index.values()).items())),
            "chance_floor_pct": 100 * sum(1 / r["n_options"] for r in index.values()) / len(index),
        },
        "runs": {}, "comparisons": [], "notes": config.get("notes", []),
    }
    scores = {}
    require(bool(config["runs"]), "At least one run is required")
    for run in config["runs"]:
        name = run["name"]
        require(isinstance(name, str) and name, "Run name must be a nonempty string")
        require(name not in scores, f"Duplicate run name: {name}")
        require(run.get("status") in ("current", "historical"), f"Run {name}: set status current or historical")
        path = base / run["path"]
        rows, validation = load_scores(path, index)
        scores[name] = rows
        report["runs"][name] = {
            "status": run["status"], "source_path": run["path"], "source_sha256": sha256(path),
            "provenance": run.get("provenance", {}), "validation": validation, "summary": summarize(rows, index),
        }
    for pair in config.get("comparisons", []):
        a, b = pair["a"], pair["b"]
        require(a in scores and b in scores and a != b, f"Invalid comparison: {a} versus {b}")
        require(bool(pair.get("classification")), f"Comparison {a} versus {b} needs a classification")
        historical = any(report["runs"][name]["status"] == "historical" for name in (a, b))
        require(not historical or "historical" in pair["classification"], f"Comparison {a} versus {b} must be labeled historical")
        report["comparisons"].append({**pair, "contains_historical_run": historical, "statistics": compare(scores[a], scores[b])})
    return report


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--output-prefix", type=Path, help="Default: analysis beside config")
    args = parser.parse_args()
    try:
        report = analyze(args.config)
    except (ValidationError, KeyError, TypeError, OSError, json.JSONDecodeError) as error:
        print(f"VALIDATION FAILED: {error}", file=sys.stderr)
        return 1
    prefix = args.output_prefix or args.config.parent / "analysis"
    prefix.parent.mkdir(parents=True, exist_ok=True)
    # Do not write a success report until every run and comparison has validated.
    json_path = prefix.with_suffix(".json")
    md_path = prefix.with_suffix(".md")
    json_path.write_text(json.dumps(report, indent=2, allow_nan=False) + "\n")
    md_path.write_text(markdown(report))
    print(f"Validated {len(report['runs'])} run(s), {report['dataset']['n']} questions each: {json_path}, {md_path}")
    return 0


if __name__ == "__main__":
    sys.exit(main())