#!/usr/bin/env python3
"""Audit and summarize the matched shared-gate vs KV-head-gate sweep."""

from __future__ import annotations

import argparse
import csv
import glob
import hashlib
import json
import subprocess
from collections import defaultdict
from pathlib import Path

import matplotlib.pyplot as plt


SUITES = ("RULER-local-full13", "BabiLong", "HELMET-ICL", "MRCR")
THRESHOLDS = (0.45, 0.50, 0.525, 0.55, 0.575, 0.60, 0.625, 0.65)
ARMS = ("token", "token_kv_head")
SLUGS = {
    "RULER-local-full13": "ruler",
    "BabiLong": "babilong",
    "HELMET-ICL": "helmet",
    "MRCR": "mrcr",
}
SCORE_DEFINITIONS = {
    "RULER-local-full13": "unweighted macro of lm-eval 8192 string-match scores over the local 13-config set",
    "BabiLong": "unweighted macro exact-match accuracy over qa1-qa5",
    "HELMET-ICL": "unweighted macro label exact-match over five ICL configurations",
    "MRCR": "unweighted macro of prefix-check plus SequenceMatcher scores over 2/4/8-needle configurations",
}


def threshold_slug(value: float) -> str:
    return {
        0.45: "045",
        0.50: "050",
        0.525: "0525",
        0.55: "055",
        0.575: "0575",
        0.60: "060",
        0.625: "0625",
        0.65: "065",
    }[value]


def one(pattern: str) -> Path:
    paths = [Path(path) for path in glob.glob(pattern, recursive=True)]
    if len(paths) != 1:
        raise RuntimeError(f"expected one path for {pattern}, found {paths}")
    return paths[0]


def read_jsonl(path: Path) -> list[dict]:
    return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]


def collect_scores(root: Path) -> dict[tuple[str, str], dict]:
    rows: dict[tuple[str, str], dict] = {}
    for split in ("a", "b"):
        result = one(str(root / f"ruler8k_splits/{split}/lm_eval/**/results_*.json"))
        payload = json.loads(result.read_text())
        for config, metrics in payload["results"].items():
            if config in payload["n-samples"]:
                rows[("RULER-local-full13", config)] = {
                    "score": float(metrics["8192,none"]),
                    "n": int(payload["n-samples"][config]["effective"]),
                }
    result = one(str(root / "babilong8k_qa1_qa5_n50/**/results_*.json"))
    payload = json.loads(result.read_text())
    for config, metrics in payload["results"].items():
        if config in payload["n-samples"]:
            rows[("BabiLong", config)] = {
                "score": float(metrics["acc,none"]),
                "n": int(payload["n-samples"][config]["effective"]),
            }
    for suite, filename in (
        ("HELMET-ICL", "helmet_icl8k_n50.jsonl"),
        ("MRCR", "mrcr_8k_2_4_8needle_n10.jsonl"),
    ):
        grouped: dict[str, list[dict]] = defaultdict(list)
        for row in read_jsonl(root / filename):
            grouped[row["config"]].append(row)
        for config, values in grouped.items():
            rows[(suite, config)] = {
                "score": sum(float(row["score"]) for row in values) / len(values),
                "n": len(values),
            }
    return rows


def sample_signatures(root: Path, score_rows: dict[tuple[str, str], dict]) -> dict[str, str]:
    signatures: dict[str, str] = {}
    for suite, rel_dirs in (
        ("RULER-local-full13", ("ruler8k_splits/a/lm_eval", "ruler8k_splits/b/lm_eval")),
        ("BabiLong", ("babilong8k_qa1_qa5_n50",)),
    ):
        configs = [config for row_suite, config in score_rows if row_suite == suite]
        for config in configs:
            matches = []
            for rel_dir in rel_dirs:
                matches.extend(root.glob(f"{rel_dir}/**/samples_{config}_*.jsonl"))
            if len(matches) != 1:
                raise RuntimeError(f"expected one sample log for {suite}/{config}: {matches}")
            values = [
                (
                    row.get("doc_id"),
                    row.get("doc_hash"),
                    row.get("prompt_hash"),
                    row.get("target_hash"),
                )
                for row in read_jsonl(matches[0])
            ]
            signatures[f"{suite}/{config}"] = hashlib.sha256(
                json.dumps(values, sort_keys=True).encode()
            ).hexdigest()
    for suite, filename in (
        ("HELMET-ICL", "helmet_icl8k_n50.jsonl"),
        ("MRCR", "mrcr_8k_2_4_8needle_n10.jsonl"),
    ):
        grouped: dict[str, list[tuple]] = defaultdict(list)
        for row in read_jsonl(root / filename):
            grouped[row["config"]].append(
                (row["row_id"], row.get("prompt_sha256"), row.get("target"))
            )
        for config, values in grouped.items():
            signatures[f"{suite}/{config}"] = hashlib.sha256(
                json.dumps(values, sort_keys=True).encode()
            ).hexdigest()
    return signatures


def read_sparsity(root: Path, suite: str) -> dict:
    if suite == "RULER-local-full13":
        parts = [
            json.loads((root / f"sparsity_ruler_{split}.json").read_text())
            for split in ("a", "b")
        ]
        sparse = sum(int(part["sparse_decisions"]) for part in parts)
        total = sum(int(part["total_decisions"]) for part in parts)
        granularities = {part.get("router_granularity") for part in parts}
        native = sum(int(part.get("native_router_decisions", 0)) for part in parts)
        phases: dict[str, dict[str, int]] = defaultdict(lambda: {"sparse": 0, "total": 0})
        for part in parts:
            for phase, values in part.get("by_phase", {}).items():
                phases[phase]["sparse"] += int(values["sparse_decisions"])
                phases[phase]["total"] += int(values["total_decisions"])
    else:
        part = json.loads((root / f"sparsity_{SLUGS[suite]}.json").read_text())
        sparse = int(part["sparse_decisions"])
        total = int(part["total_decisions"])
        granularities = {part.get("router_granularity")}
        native = int(part.get("native_router_decisions", 0))
        phases = {
            phase: {
                "sparse": int(values["sparse_decisions"]),
                "total": int(values["total_decisions"]),
            }
            for phase, values in part.get("by_phase", {}).items()
        }
    return {
        "sparse_decisions": sparse,
        "total_decisions": total,
        "sparsity": sparse / total,
        "full_attention_usage": 1.0 - sparse / total,
        "router_granularities": sorted(value for value in granularities if value),
        "native_router_decisions": native or None,
        "by_phase": phases,
    }


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


def command_output(command: list[str], cwd: Path) -> str:
    try:
        return subprocess.run(
            command, cwd=cwd, check=True, text=True, capture_output=True
        ).stdout.strip()
    except (OSError, subprocess.CalledProcessError):
        return "unavailable in downloaded HF snapshot"


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--repo", type=Path, default=Path("/data/sjm/AHA/AHA-Qwen3")
    )
    parser.add_argument(
        "--eval-root",
        type=Path,
        default=Path("/data/sjm/AHA/AHA-Qwen3/experiments/qwen3_1p7b_router_granularity_20260714/eval"),
    )
    parser.add_argument(
        "--baseline-root",
        type=Path,
        default=Path("/data/sjm/AHA/AHA-Qwen3/experiments/qwen3_1p7b_fourbench_8k_broad_20260713/vanilla"),
    )
    parser.add_argument("--output-dir", type=Path)
    parser.add_argument("--input-dir", type=Path)
    args = parser.parse_args()
    output = args.output_dir or args.eval_root.parent / "summary"
    output.mkdir(parents=True, exist_ok=True)

    baseline = collect_scores(args.baseline_root)
    baseline_signatures = sample_signatures(args.baseline_root, baseline)
    curve_rows = []
    per_config = []
    points = []
    alignment = {}
    done_files = ("ruler_a.DONE", "ruler_b.DONE", "babilong.DONE", "helmet.DONE", "mrcr.DONE")

    for arm in ARMS:
        for threshold in THRESHOLDS:
            method = f"{arm}_t{threshold_slug(threshold)}"
            root = args.eval_root / method
            missing = [name for name in done_files if not (root / name).exists()]
            if missing:
                raise RuntimeError(f"incomplete {method}: missing {missing}")
            scores = collect_scores(root)
            if set(scores) != set(baseline):
                raise RuntimeError(f"configuration mismatch for {method}")
            if any(scores[key]["n"] != baseline[key]["n"] for key in baseline):
                raise RuntimeError(f"sample-count mismatch for {method}")
            signatures = sample_signatures(root, scores)
            mismatches = sorted(
                key for key in baseline_signatures if signatures.get(key) != baseline_signatures[key]
            )
            alignment[method] = {"aligned": not mismatches, "mismatches": mismatches}
            if mismatches:
                raise RuntimeError(f"prompt/target hash mismatch for {method}: {mismatches}")

            suite_rows = []
            sparse_sum = total_sum = 0
            for suite in SUITES:
                keys = sorted(key for key in baseline if key[0] == suite)
                baseline_score = sum(baseline[key]["score"] for key in keys) / len(keys)
                score = sum(scores[key]["score"] for key in keys) / len(keys)
                sparsity = read_sparsity(root, suite)
                expected_granularity = arm
                if sparsity["router_granularities"] != [expected_granularity]:
                    raise RuntimeError(
                        f"{method}/{suite} recorded {sparsity['router_granularities']}, "
                        f"expected {[expected_granularity]}"
                    )
                sparse_sum += sparsity["sparse_decisions"]
                total_sum += sparsity["total_decisions"]
                row = {
                    "arm": arm,
                    "router_granularity": arm,
                    "threshold": threshold,
                    "benchmark": suite,
                    "config_count": len(keys),
                    "samples_per_config": ",".join(
                        str(value) for value in sorted({baseline[key]["n"] for key in keys})
                    ),
                    "total_samples": sum(baseline[key]["n"] for key in keys),
                    "tuned_vanilla_score": baseline_score,
                    "quality_floor_95pct": 0.95 * baseline_score,
                    "score": score,
                    "retention": score / baseline_score if baseline_score else None,
                    "quality_pass": score >= 0.95 * baseline_score,
                    **{key: sparsity[key] for key in (
                        "sparse_decisions", "total_decisions", "sparsity",
                        "full_attention_usage", "native_router_decisions", "by_phase"
                    )},
                }
                curve_rows.append(row)
                suite_rows.append(row)
                for key in keys:
                    per_config.append({
                        "arm": arm,
                        "threshold": threshold,
                        "benchmark": suite,
                        "config": key[1],
                        "n": baseline[key]["n"],
                        "tuned_vanilla_score": baseline[key]["score"],
                        "score": scores[key]["score"],
                    })
            points.append({
                "arm": arm,
                "threshold": threshold,
                "all_suite_quality_pass": all(row["quality_pass"] for row in suite_rows),
                "decision_weighted_sparsity": sparse_sum / total_sum,
                "full_attention_usage": 1.0 - sparse_sum / total_sum,
                "sparse_decisions": sparse_sum,
                "total_decisions": total_sum,
                "suites": suite_rows,
            })

    headline = {}
    for arm in ARMS:
        valid = [
            point for point in points
            if point["arm"] == arm and point["all_suite_quality_pass"]
        ]
        if valid:
            best = max(valid, key=lambda point: point["decision_weighted_sparsity"])
            headline[arm] = {"found": True, **best}
        else:
            headline[arm] = {
                "found": False,
                "statement": "No measured threshold preserved at least 95% of tuned vanilla on every suite.",
            }

    input_dir = args.input_dir or args.baseline_root.parent / "inputs"
    frozen_inputs = [
        input_dir / "helmet_icl_8k_n50_per_config.jsonl",
        input_dir / "mrcr_8k_2_4_8needle_n10_per_config.jsonl",
    ]
    payload = {
        "protocol": {
            "model": "Qwen3-1.7B tuned vanilla",
            "context_length": 8192,
            "arms": {
                "token": "L2A-style shared-gate: one native gate per token/layer",
                "token_kv_head": "AHA: one native gate per token/KV-head/layer",
            },
            "local_attention": {"sink_tokens": 64, "recent_tokens": 256},
            "inference": "strict AHA routing in both prefill and decode; no force-full heads or full-decode fallback",
            "thresholds": list(THRESHOLDS),
            "quality_rule": "every suite macro score >= 95% of aligned tuned-vanilla macro score",
            "headline_rule": "highest decision-weighted measured sparsity among thresholds passing every suite",
            "sparsity_definition": "hard local routes / token x KV-head x layer effective decisions",
            "suite_scope": "RULER local full13 means this repository's fixed 13-config set, not complete upstream RULER",
        },
        "score_definitions": SCORE_DEFINITIONS,
        "headline": headline,
        "points": points,
        "curves": curve_rows,
        "per_config": per_config,
        "alignment": {"all_aligned": all(value["aligned"] for value in alignment.values()), "methods": alignment},
        "appendix_note": (
            "Upstream L2A results are not merged into this table because architecture, "
            "training data, objective, and evaluation protocols differ."
        ),
        "reproducibility": {
            "git_commit": command_output(["git", "rev-parse", "HEAD"], args.repo),
            "git_status_short": command_output(["git", "status", "--short"], args.repo),
            "container": "dmtd-repro",
            "container_image": command_output(
                ["docker", "inspect", "-f", "{{.Config.Image}}", "dmtd-repro"], args.repo
            ),
            "baseline_root": str(args.baseline_root),
            "eval_root": str(args.eval_root),
            "frozen_input_sha256": {
                str(path): sha256(path) for path in frozen_inputs if path.exists()
            },
        },
    }
    (output / "router_granularity_results.json").write_text(
        json.dumps(payload, indent=2) + "\n"
    )

    flat_curve_rows = [
        {key: value for key, value in row.items() if key != "by_phase"}
        for row in curve_rows
    ]
    with (output / "quality_sparsity_curves.csv").open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(flat_curve_rows[0]))
        writer.writeheader()
        writer.writerows(flat_curve_rows)
    with (output / "per_config_scores.csv").open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(per_config[0]))
        writer.writeheader()
        writer.writerows(per_config)

    figure, axes = plt.subplots(2, 2, figsize=(12, 9), constrained_layout=True)
    colors = {"token": "#d62728", "token_kv_head": "#1f77b4"}
    labels = {"token": "L2A-style shared-gate", "token_kv_head": "AHA KV-head gate"}
    for axis, suite in zip(axes.flat, SUITES):
        for arm in ARMS:
            values = sorted(
                (row for row in curve_rows if row["arm"] == arm and row["benchmark"] == suite),
                key=lambda row: row["sparsity"],
            )
            axis.plot(
                [100 * row["sparsity"] for row in values],
                [100 * row["retention"] for row in values],
                marker="o", color=colors[arm], label=labels[arm],
            )
            for row in values:
                axis.annotate(f"{row['threshold']:.3g}", (100 * row["sparsity"], 100 * row["retention"]), fontsize=7)
        axis.axhline(95, color="black", linestyle="--", linewidth=1)
        axis.set_title(suite)
        axis.set_xlabel("Measured effective sparsity (%)")
        axis.set_ylabel("Retention vs tuned vanilla (%)")
        axis.grid(alpha=0.25)
    axes.flat[0].legend()
    figure.suptitle("Qwen3-1.7B matched router-granularity quality–sparsity curves")
    figure.savefig(output / "quality_sparsity_curves.png", dpi=180)
    plt.close(figure)
    print(json.dumps(headline, indent=2))


if __name__ == "__main__":
    main()