#!/usr/bin/env python
from __future__ import annotations

import argparse
import json
import math
import subprocess
import sys
from pathlib import Path
from typing import Any

PROJECT_ROOT = Path(__file__).resolve().parents[1]


DEFAULT_RUNS = [
    ("local_atlas", [Path("runs/local_atlas_val_proxy")]),
    ("task_memory", [Path("runs/task_memory_val_proxy")]),
    ("ctt_residual_full", sorted(Path("runs").glob("ctt_residual_full_seed*_val_proxy"))),
    (
        "ctt_residual_base_context",
        sorted(Path("runs").glob("ctt_residual_base_context_seed*_val_proxy")),
    ),
    (
        "ctt_residual_base_context_obs",
        sorted(Path("runs").glob("ctt_residual_base_context_obs_seed*_val_proxy")),
    ),
    (
        "ctt_residual_base_context_obj",
        sorted(Path("runs").glob("ctt_residual_base_context_obj_seed*_val_proxy")),
    ),
    (
        "ctt_residual_base_context_obs_obj",
        sorted(Path("runs").glob("ctt_residual_base_context_obs_obj_seed*_val_proxy")),
    ),
    (
        "ctt_gated_residual_full",
        sorted(Path("runs").glob("ctt_gated_residual_full_seed*_val_proxy")),
    ),
]


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Build CTT validation proxy comparison table.")
    parser.add_argument("--out-dir", type=Path, default=Path("runs/ctt_val_proxy_comparison"))
    parser.add_argument("--safety-slack", type=float, default=0.01)
    parser.add_argument(
        "--no-markdown-report",
        action="store_true",
        help="Do not write report.md; the persistent prose summary lives in README.md.",
    )
    args = parser.parse_args(argv)

    rows = [
        _row(name, run_dirs)
        for name, run_dirs in DEFAULT_RUNS
        if run_dirs and all((run_dir / "metrics.json").exists() for run_dir in run_dirs)
    ]
    if not rows:
        raise SystemExit("no proxy runs found")
    baseline = next((row for row in rows if row["method"] == "local_atlas"), None)
    if baseline is None:
        raise SystemExit("local_atlas baseline is required for gate comparison")
    for row in rows:
        row["safety_ok_vs_local_atlas"] = (
            row["negative_near_0p20"] <= baseline["negative_near_0p20"] + args.safety_slack
        )
        row["beats_local_atlas_support"] = (
            row["pptc_0p20"] > baseline["pptc_0p20"]
            or row["pptc_0p40"] > baseline["pptc_0p40"]
            or row["mean_positive_distance"] < baseline["mean_positive_distance"]
        )
        row["proxy_gate_pass"] = bool(
            row["method"].startswith("ctt")
            and row["safety_ok_vs_local_atlas"]
            and row["beats_local_atlas_support"]
        )

    out_dir = args.out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    payload = {
        "report_type": "ctt_val_proxy_comparison",
        "schema_version": 2,
        "baseline": "local_atlas",
        "safety_slack": args.safety_slack,
        "rows": rows,
        "data_hash": _first(rows, "data_hash"),
        "split_hash": _first(rows, "split_hash"),
    }
    (out_dir / "metrics.json").write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    (out_dir / "metrics_by_task.json").write_text(
        json.dumps(_grouped_metrics(DEFAULT_RUNS, group_key="task_id"), indent=2, sort_keys=True)
        + "\n"
    )
    (out_dir / "metrics_by_seed.json").write_text(
        json.dumps(_grouped_metrics(DEFAULT_RUNS, group_key="seed"), indent=2, sort_keys=True)
        + "\n"
    )
    (out_dir / "table.tex").write_text(_table(rows) + "\n")
    _write_report_artifact(
        out_dir,
        rows,
        safety_slack=args.safety_slack,
        no_markdown_report=args.no_markdown_report,
    )
    (out_dir / "train.log").write_text("comparison artifact; no training\n")
    (out_dir / "eval.log").write_text(
        "\n".join(f"{row['method']}: {row['run_path']}" for row in rows) + "\n"
    )
    (out_dir / "config.yaml").write_text(f"safety_slack: {args.safety_slack}\n")
    (out_dir / "command.txt").write_text(
        "python scripts/build_ctt_proxy_comparison.py " + " ".join(sys.argv[1:]) + "\n"
    )
    (out_dir / "git_hash.txt").write_text(_run(["git", "rev-parse", "HEAD"]) + "\n")
    (out_dir / "data_hash.txt").write_text(str(payload["data_hash"] or "") + "\n")
    (out_dir / "split_hash.txt").write_text(str(payload["split_hash"] or "") + "\n")
    print(json.dumps({"out_dir": str(out_dir), "rows": len(rows)}, indent=2))
    return 0


def _row(method: str, run_dirs: list[Path]) -> dict[str, Any]:
    payloads = [json.loads((run_dir / "metrics.json").read_text()) for run_dir in run_dirs]
    summaries = [payload.get("summary", {}) for payload in payloads]
    return {
        "method": method,
        "run_path": ",".join(str(run_dir) for run_dir in run_dirs),
        "num_rows": sum(int(payload.get("num_rows", 0)) for payload in payloads),
        "train_seeds": len(run_dirs),
        "pptc_0p20": _mean_across(summaries, "pptc_at_16_thr_0p20"),
        "pptc_0p40": _mean_across(summaries, "pptc_at_16_thr_0p40"),
        "negative_near_0p20": _mean_across(summaries, "negative_near_at_16_thr_0p20"),
        "negative_near_0p40": _mean_across(summaries, "negative_near_at_16_thr_0p40"),
        "mean_positive_distance": _mean_across(summaries, "mean_positive_distance_at_16"),
        "mean_negative_distance": _mean_across(summaries, "mean_negative_distance_at_16"),
        "pos_closer_than_neg": _mean_across(summaries, "pos_closer_than_neg_at_16"),
        "candidate_diversity": _mean_across(summaries, "candidate_diversity_at_16"),
        "collapse_rate": _mean_across(summaries, "collapse_rate_at_16"),
        "proxy_support_distance": _mean_across(summaries, "proxy_support_distance_at_16"),
        "data_hash": payloads[0].get("data_hash"),
        "split_hash": payloads[0].get("target_split_hash", payloads[0].get("split_hash")),
    }


def _mean_across(summaries: list[dict[str, Any]], key: str) -> float:
    values = []
    for summary in summaries:
        value = summary.get(key, {}).get("micro", {}).get("mean")
        if isinstance(value, (int, float)) and math.isfinite(float(value)):
            values.append(float(value))
    return sum(values) / len(values) if values else math.nan


def _table(rows: list[dict[str, Any]]) -> str:
    lines = [
        "% Auto-generated by scripts/build_ctt_proxy_comparison.py",
        "\\begin{tabular}{lrrrrrrrrrrc}",
        "\\toprule",
        "Method & Seeds & Rows & PPTC@0.20 & PPTC@0.40 & Neg@0.20 & Neg@0.40 & Pos<Neg & MeanPos & MeanNeg & Collapse & Gate \\\\",
        "\\midrule",
    ]
    for row in rows:
        lines.append(
            f"{_latex_escape(row['method'])} & {row['train_seeds']} & {row['num_rows']} & "
            f"{_fmt(row['pptc_0p20'])} & {_fmt(row['pptc_0p40'])} & "
            f"{_fmt(row['negative_near_0p20'])} & {_fmt(row['negative_near_0p40'])} & "
            f"{_fmt(row['pos_closer_than_neg'])} & {_fmt(row['mean_positive_distance'])} & "
            f"{_fmt(row['mean_negative_distance'])} & {_fmt(row['collapse_rate'])} & "
            f"{_gate(row)} \\\\"
        )
    lines.extend(["\\bottomrule", "\\end{tabular}"])
    return "\n".join(lines)


def _report(rows: list[dict[str, Any]], safety_slack: float) -> str:
    lines = [
        "# CTT Validation Proxy Comparison",
        "",
        f"Safety slack over local-atlas NegativeNear@0.20: `{safety_slack:.4f}`",
        "",
        "| Method | Seeds | Rows | PPTC@0.20 | PPTC@0.40 | Neg@0.20 | Neg@0.40 | Pos<Neg | MeanPos | MeanNeg | Diversity | Collapse | Gate | Run |",
        "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- | --- |",
    ]
    for row in rows:
        lines.append(
            f"| {row['method']} | {row['train_seeds']} | {row['num_rows']} | "
            f"{_fmt(row['pptc_0p20'])} | "
            f"{_fmt(row['pptc_0p40'])} | {_fmt(row['negative_near_0p20'])} | "
            f"{_fmt(row['negative_near_0p40'])} | {_fmt(row['pos_closer_than_neg'])} | "
            f"{_fmt(row['mean_positive_distance'])} | {_fmt(row['mean_negative_distance'])} | "
            f"{_fmt(row['candidate_diversity'])} | {_fmt(row['collapse_rate'])} | "
            f"{_gate(row)} | `{row['run_path']}` |"
        )
    lines.append("")
    lines.append("Gate is proxy-only. It is not OutcomePTR or measured rollout success.")
    return "\n".join(lines)


def _write_report_artifact(
    out_dir: Path,
    rows: list[dict[str, Any]],
    *,
    safety_slack: float,
    no_markdown_report: bool,
) -> None:
    report_path = out_dir / "report.md"
    if no_markdown_report:
        if report_path.exists():
            report_path.unlink()
        return
    report_path.write_text(_report(rows, safety_slack) + "\n")


def _grouped_metrics(
    run_specs: list[tuple[str, list[Path]]],
    *,
    group_key: str,
) -> dict[str, dict[str, dict[str, float]]]:
    output: dict[str, dict[str, dict[str, float]]] = {}
    for method, run_dirs in run_specs:
        if not run_dirs or not all((run_dir / "metrics.json").exists() for run_dir in run_dirs):
            continue
        rows: list[dict[str, Any]] = []
        for run_dir in run_dirs:
            payload = json.loads((run_dir / "metrics.json").read_text())
            rows.extend(payload.get("rows", []))
        if not rows:
            continue
        method_groups: dict[str, dict[str, list[float]]] = {}
        for row in rows:
            group = str(row.get(group_key, "unknown"))
            metrics = method_groups.setdefault(group, {})
            for source_key, target_key in _row_metric_key_map().items():
                value = row.get(source_key)
                if isinstance(value, (int, float)) and math.isfinite(float(value)):
                    metrics.setdefault(target_key, []).append(float(value))
        output[method] = {
            group: {name: sum(values) / len(values) for name, values in sorted(metrics.items())}
            for group, metrics in sorted(method_groups.items())
        }
    return output


def _row_metric_key_map() -> dict[str, str]:
    return {
        "pptc_at_16_thr_0p20": "pptc_0p20",
        "pptc_at_16_thr_0p40": "pptc_0p40",
        "negative_near_at_16_thr_0p20": "negative_near_0p20",
        "negative_near_at_16_thr_0p40": "negative_near_0p40",
        "mean_positive_distance_at_16": "mean_positive_distance",
        "mean_negative_distance_at_16": "mean_negative_distance",
        "pos_closer_than_neg_at_16": "pos_closer_than_neg",
        "candidate_diversity_at_16": "candidate_diversity",
        "collapse_rate_at_16": "collapse_rate",
        "proxy_support_distance_at_16": "proxy_support_distance",
    }


def _fmt(value: Any) -> str:
    if not isinstance(value, (int, float)) or not math.isfinite(float(value)):
        return "n/a"
    return f"{float(value):.4f}"


def _gate(row: dict[str, Any]) -> str:
    if row.get("proxy_gate_pass"):
        return "pass"
    if row["method"].startswith("ctt"):
        return "fail"
    return "baseline"


def _latex_escape(value: str) -> str:
    return value.replace("_", "\\_").replace("%", "\\%").replace("&", "\\&")


def _first(rows: list[dict[str, Any]], key: str) -> Any:
    for row in rows:
        value = row.get(key)
        if value:
            return value
    return None


def _run(command: list[str]) -> str:
    try:
        return subprocess.check_output(command, cwd=PROJECT_ROOT, text=True).strip()
    except (subprocess.CalledProcessError, FileNotFoundError):
        return ""


if __name__ == "__main__":
    raise SystemExit(main())