#!/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]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Fit a conformal-style lower confidence bound for causal dominance margins."
    )
    parser.add_argument("--input", type=Path, required=True, help="Measured calibration JSON rows.")
    parser.add_argument("--out-dir", type=Path, default=Path("runs/dominance_calibration"))
    parser.add_argument("--alpha", type=float, default=0.1)
    parser.add_argument(
        "--no-markdown-report",
        action="store_true",
        help="Do not write report.md; persistent prose is consolidated in README.md.",
    )
    args = parser.parse_args(argv)
    if not 0.0 < args.alpha < 1.0:
        parser.error("--alpha must be in (0, 1)")

    rows = _rows(json.loads(args.input.read_text()))
    residuals = []
    for index, row in enumerate(rows):
        if "predicted_margin" in row and "measured_margin" in row:
            predicted = float(row["predicted_margin"])
            measured = float(row["measured_margin"])
        elif "predicted_scores" in row and "utilities" in row and "base_index" in row and "selected_index" in row:
            predicted_scores = [float(value) for value in row["predicted_scores"]]
            utilities = [float(value) for value in row["utilities"]]
            base_index = int(row["base_index"])
            selected_index = int(row["selected_index"])
            predicted = predicted_scores[selected_index] - predicted_scores[base_index]
            measured = utilities[selected_index] - utilities[base_index]
        else:
            raise SystemExit(
                f"row {index} needs predicted_margin/measured_margin or "
                "predicted_scores/utilities/base_index/selected_index"
            )
        residuals.append(abs(measured - predicted))
    if not residuals:
        raise SystemExit("calibration input has no measured rows")

    residuals.sort()
    q_index = min(len(residuals) - 1, math.ceil((1.0 - args.alpha) * (len(residuals) + 1)) - 1)
    quantile = residuals[max(0, q_index)]
    out_dir = args.out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    _write_provenance(out_dir, args)
    payload = {
        "report_type": "dominance_calibration",
        "alpha": args.alpha,
        "num_rows": len(residuals),
        "residual_quantile": quantile,
        "rule": "execute candidate only if predicted_margin - residual_quantile > tau",
    }
    (out_dir / "metrics.json").write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    (out_dir / "metrics_by_task.json").write_text("{}\n")
    (out_dir / "metrics_by_seed.json").write_text("{}\n")
    (out_dir / "train.log").write_text("fit conformal residual quantile from measured calibration rows\n")
    (out_dir / "eval.log").write_text("no held-out dominance eval in calibration script\n")
    (out_dir / "table.tex").write_text(_table(payload) + "\n")
    _write_markdown_report(out_dir, payload, no_markdown_report=args.no_markdown_report)
    print(json.dumps({"out_dir": str(out_dir), "residual_quantile": quantile}, indent=2))
    return 0


def _rows(payload: Any) -> list[dict[str, Any]]:
    rows = payload.get("rows", payload) if isinstance(payload, dict) else payload
    if not isinstance(rows, list):
        raise SystemExit("input must be a JSON list or object with rows")
    return rows


def _write_provenance(out_dir: Path, args: argparse.Namespace) -> None:
    (out_dir / "config.yaml").write_text(
        "\n".join(f"{key}: {value}" for key, value in sorted(vars(args).items())) + "\n"
    )
    (out_dir / "command.txt").write_text("python scripts/calibrate_dominance.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("input_json:" + _sha256(args.input) + "\n")
    (out_dir / "split_hash.txt").write_text("calibration_input\n")


def _sha256(path: Path) -> str:
    import hashlib

    h = hashlib.sha256()
    h.update(path.read_bytes())
    return h.hexdigest()


def _table(payload: dict[str, Any]) -> str:
    return "\n".join(
        [
            "% Auto-generated by scripts/calibrate_dominance.py",
            "\\begin{tabular}{lrr}",
            "\\toprule",
            "Rows & Alpha & Residual quantile \\\\",
            "\\midrule",
            f"{payload['num_rows']} & {payload['alpha']:.3f} & {payload['residual_quantile']:.4f} \\\\",
            "\\bottomrule",
            "\\end{tabular}",
        ]
    )


def _report(payload: dict[str, Any]) -> str:
    return "\n".join(
        [
            "# Dominance Calibration",
            "",
            f"Rows: `{payload['num_rows']}`",
            f"Alpha: `{payload['alpha']}`",
            f"Residual quantile: `{payload['residual_quantile']:.6f}`",
            "",
            f"Rule: `{payload['rule']}`",
        ]
    )


def _write_markdown_report(
    out_dir: Path,
    payload: dict[str, Any],
    *,
    no_markdown_report: bool,
) -> None:
    report_path = out_dir / "report.md"
    if no_markdown_report:
        report_path.unlink(missing_ok=True)
        return
    report_path.write_text(_report(payload) + "\n")


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())