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

import argparse
import glob
import hashlib
import json
import math
import subprocess
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any


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


DEFAULT_PATTERNS = (
    "runs/ctt_base_context_obs_learned_dominance_chartcompat_obs_utility_task_envclip_k16_train_to_test/metrics.json",
    "runs/ctt_base_context_obs_learned_dominance_*bundle*_envclip_k16_train_to_test/metrics.json",
    "runs/ctt_base_context_obs_dominance_envclip_k16_train_to_test/metrics.json",
    "runs/ctt_base_context_obs_dominance_envclip_k16_train_to_test_tau0/metrics.json",
    "runs/ctt_dominance_utility_energy_val_to_test_seed*/metrics.json",
    "runs/ctt_base_context_obs_learned_dominance_*_tanh_train_to_test/metrics.json",
    "runs/ctt_base_context_obs_dominance_tanh_train_to_test/metrics.json",
    "runs/ctt_base_context_obs_learned_dominance_*_perdim_trainmax_train_to_test/metrics.json",
    "runs/ctt_base_context_obs_dominance_perdim_trainmax_train_to_test/metrics.json",
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Build a non-cherry-picked selector diagnostic sweep table from "
            "completed CTT selector metrics.json files."
        )
    )
    parser.add_argument(
        "--metrics",
        action="append",
        default=[],
        help="Metrics file or glob. Defaults cover current env_clip/tanh/per-dim selector runs.",
    )
    parser.add_argument("--out-dir", type=Path, default=Path("runs/ctt_selector_diagnostic_sweep"))
    parser.add_argument("--selected-min", type=float, default=0.4745)
    parser.add_argument("--proposal-oracle-min", type=float, default=0.50)
    parser.add_argument("--selector-gap-max", type=float, default=0.03)
    args = parser.parse_args(argv)

    metric_paths = _resolve_metric_paths(args.metrics or list(DEFAULT_PATTERNS))
    if not metric_paths:
        raise SystemExit("no selector metrics found")

    rows = [_row(path) for path in metric_paths]
    best_rows = _best_by_family(rows)
    gates = [_gate(row, args) for row in best_rows]

    out_dir = args.out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    payload = {
        "report_type": "ctt_selector_diagnostic_sweep",
        "schema_version": 1,
        "selection_rule": "best selected_success per diagnostic family; all candidate rows retained",
        "thresholds": {
            "selected_min": args.selected_min,
            "proposal_oracle_min": args.proposal_oracle_min,
            "selector_gap_max": args.selector_gap_max,
        },
        "num_inputs": len(metric_paths),
        "input_metrics": [str(path) for path in metric_paths],
        "rows": rows,
        "best_by_family": best_rows,
        "gates": gates,
        "overall_pass": all(gate["pass"] for gate in gates),
        "data_hash": _combined_hash([row.get("data_hash") for row in rows]),
        "split_hash": _combined_hash([row.get("split_hash") for row in rows]),
    }
    (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(_group_rows(rows, "family"), indent=2, sort_keys=True) + "\n"
    )
    (out_dir / "metrics_by_seed.json").write_text(
        json.dumps(_group_rows(rows, "seed"), indent=2, sort_keys=True) + "\n"
    )
    (out_dir / "table.tex").write_text(_table(best_rows) + "\n")
    (out_dir / "config.yaml").write_text(_config(args, metric_paths) + "\n")
    (out_dir / "command.txt").write_text(
        "python scripts/build_selector_diagnostic_sweep.py " + " ".join(sys.argv[1:]) + "\n"
    )
    (out_dir / "git_hash.txt").write_text(_git_hash() + "\n")
    (out_dir / "data_hash.txt").write_text(str(payload["data_hash"]) + "\n")
    (out_dir / "split_hash.txt").write_text(str(payload["split_hash"]) + "\n")
    (out_dir / "train.log").write_text("selector sweep artifact; source selectors trained separately\n")
    (out_dir / "eval.log").write_text(
        "\n".join(
            [
                f"num_inputs={len(metric_paths)}",
                f"families={','.join(row['family'] for row in best_rows)}",
                f"overall_pass={payload['overall_pass']}",
            ]
        )
        + "\n"
    )
    print(
        json.dumps(
            {
                "out_dir": str(out_dir),
                "num_inputs": len(metric_paths),
                "families": [row["family"] for row in best_rows],
                "overall_pass": payload["overall_pass"],
            },
            indent=2,
        )
    )
    return 0


def _resolve_metric_paths(patterns: list[str]) -> list[Path]:
    paths: list[Path] = []
    for pattern in patterns:
        if any(char in pattern for char in "*?[]"):
            matches = [Path(item) for item in sorted(glob.glob(pattern))]
        else:
            matches = [Path(pattern)]
        for path in matches:
            if path.exists() and path.name == "metrics.json" and path not in paths:
                paths.append(path)
    return paths


def _row(path: Path) -> dict[str, Any]:
    data = json.loads(path.read_text())
    summary = data.get("eval_summary") or _micro_summary(data.get("summary", {}))
    run_name = path.parent.name
    family = _family(run_name, data)
    selector = _selector_name(run_name, data)
    return {
        "run_path": str(path.parent),
        "family": family,
        "selector": selector,
        "seed": _infer_seed(run_name, data),
        "report_type": data.get("report_type", "unknown"),
        "k": int(data.get("k") or _infer_k(run_name)),
        "base_success": _num(summary.get("base_success")),
        "selected_success": _num(summary.get("selected_success")),
        "proposal_oracle_success": _num(summary.get("proposal_oracle_success")),
        "hidden_chart_oracle_success": _num(summary.get("hidden_chart_oracle_success")),
        "coverage": _num(summary.get("coverage")),
        "fallback_rate": _num(summary.get("fallback_rate")),
        "success_support_gap": _num(summary.get("success_support_gap")),
        "success_selector_gap": _num(summary.get("success_selector_gap")),
        "outcome_ptr": _num(summary.get("outcome_ptr")),
        "calibration_ece": _num(summary.get("pairwise_causal_calibration_ece")),
        "selector_regret": _num(summary.get("selector_regret")),
        "data_hash": _first_hash(data, ("data_hash", "eval_target_content_hash", "selector_eval_target_content_hash")),
        "split_hash": _first_hash(data, ("split_hash", "eval_target_split_hash", "selector_eval_target_split_hash")),
    }


def _micro_summary(summary: dict[str, Any]) -> dict[str, Any]:
    output: dict[str, Any] = {}
    for name, payload in summary.items():
        if isinstance(payload, dict):
            output[name] = payload.get("micro", {}).get("mean")
    return output


def _family(run_name: str, data: dict[str, Any]) -> str:
    k = int(data.get("k") or _infer_k(run_name))
    if "ctt_dominance_utility_energy" in run_name or "utility_energy" in run_name:
        return f"K{k} env_clip utility-energy"
    if "envclip_k16" in run_name:
        return "K16 env_clip"
    if "envclip" in run_name:
        return f"K{k} env_clip"
    if "tanh" in run_name:
        return f"K{k} tanh"
    if "perdim_trainmax" in run_name:
        return f"K{k} per-dim trainmax"
    return f"K{k} other"


def _selector_name(run_name: str, data: dict[str, Any]) -> str:
    report_type = str(data.get("report_type", ""))
    if data.get("score_source") == "checkpoint" or "utility_energy" in run_name:
        tau_mode = str(data.get("tau_mode", "auto"))
        return f"checkpoint utility energy/LCB {tau_mode}, seed={_infer_seed(run_name, data)}"
    if report_type == "dominance_calibrated_selector_eval":
        tau_mode = str(data.get("tau_mode", "auto"))
        return f"LCB {tau_mode}"
    feature_set = str(data.get("feature_set", "unknown"))
    target = str(data.get("target", "unknown"))
    extras = []
    if data.get("success_bonus") not in {None, 0, 0.0}:
        extras.append(f"bonus={data['success_bonus']}")
    if "chartcompat_obs" in run_name and "chartcompat" not in feature_set:
        extras.append("chartcompat_obs")
    suffix = ", " + ", ".join(extras) if extras else ""
    return f"{feature_set}/{target}{suffix}"


def _infer_k(run_name: str) -> int:
    return 16 if "k16" in run_name else 8


def _infer_seed(run_name: str, data: dict[str, Any]) -> str:
    seed = data.get("seed")
    if seed is not None:
        return str(seed)
    marker = "seed"
    if marker in run_name:
        suffix = run_name.rsplit(marker, 1)[-1]
        digits = []
        for char in suffix:
            if char.isdigit():
                digits.append(char)
            else:
                break
        if digits:
            return "".join(digits)
    return "pooled"


def _best_by_family(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        grouped[row["family"]].append(row)
    best = []
    for family, items in grouped.items():
        best.append(
            max(
                items,
                key=lambda row: (
                    _sort_num(row.get("selected_success")),
                    _sort_num(row.get("proposal_oracle_success")),
                    -_sort_num(row.get("success_selector_gap")),
                ),
            )
        )
    return sorted(best, key=lambda row: (row["k"], row["family"]))


def _gate(row: dict[str, Any], args: argparse.Namespace) -> dict[str, Any]:
    selected = _num(row.get("selected_success"))
    proposal = _num(row.get("proposal_oracle_success"))
    selector_gap = _num(row.get("success_selector_gap"))
    passed = (
        selected is not None
        and proposal is not None
        and selector_gap is not None
        and selected >= args.selected_min
        and proposal >= args.proposal_oracle_min
        and selector_gap <= args.selector_gap_max
    )
    return {
        "family": row["family"],
        "selector": row["selector"],
        "pass": bool(passed),
        "status": "method_success" if passed else "diagnostic_only",
        "selected_success": selected,
        "proposal_oracle_success": proposal,
        "success_selector_gap": selector_gap,
    }


def _group_rows(rows: list[dict[str, Any]], group_key: str) -> dict[str, dict[str, float]]:
    metrics = (
        "base_success",
        "selected_success",
        "proposal_oracle_success",
        "coverage",
        "success_support_gap",
        "success_selector_gap",
        "outcome_ptr",
        "calibration_ece",
    )
    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        grouped[str(row.get(group_key, "unknown"))].append(row)
    output: dict[str, dict[str, float]] = {}
    for group, items in sorted(grouped.items()):
        output[group] = {}
        for metric in metrics:
            values = [_num(item.get(metric)) for item in items]
            clean = [value for value in values if value is not None]
            if clean:
                output[group][metric] = sum(clean) / len(clean)
    return output


def _table(rows: list[dict[str, Any]]) -> str:
    lines = [
        "% Auto-generated by scripts/build_selector_diagnostic_sweep.py",
        "\\begin{tabular}{llrrrrrr}",
        "\\toprule",
        "Family & Best selector & Base & Selected & Proposal & Coverage & Sel. gap & Support gap \\\\",
        "\\midrule",
    ]
    for row in rows:
        lines.append(
            f"{_latex(row['family'])} & {_latex(row['selector'])} & "
            f"{_fmt(row.get('base_success'))} & {_fmt(row.get('selected_success'))} & "
            f"{_fmt(row.get('proposal_oracle_success'))} & {_fmt(row.get('coverage'))} & "
            f"{_fmt(row.get('success_selector_gap'))} & {_fmt(row.get('success_support_gap'))} \\\\"
        )
    lines.extend(["\\bottomrule", "\\end{tabular}"])
    return "\n".join(lines)


def _config(args: argparse.Namespace, paths: list[Path]) -> str:
    return "\n".join(
        [
            f"out_dir: {args.out_dir}",
            f"selected_min: {args.selected_min}",
            f"proposal_oracle_min: {args.proposal_oracle_min}",
            f"selector_gap_max: {args.selector_gap_max}",
            "metrics:",
            *[f"  - {path}" for path in paths],
        ]
    )


def _first_hash(data: dict[str, Any], keys: tuple[str, ...]) -> str | None:
    for key in keys:
        value = data.get(key)
        if isinstance(value, str) and value:
            return value
    return None


def _combined_hash(values: list[Any]) -> str:
    clean = [str(value) for value in values if value not in {None, ""}]
    blob = json.dumps(sorted(clean), separators=(",", ":")).encode()
    return hashlib.sha256(blob).hexdigest()


def _git_hash() -> str:
    try:
        return subprocess.check_output(
            ["git", "rev-parse", "HEAD"],
            cwd=PROJECT_ROOT,
            text=True,
            stderr=subprocess.DEVNULL,
        ).strip()
    except (OSError, subprocess.CalledProcessError):
        return "unknown"


def _num(value: Any) -> float | None:
    if value is None:
        return None
    try:
        numeric = float(value)
    except (TypeError, ValueError):
        return None
    return numeric if math.isfinite(numeric) else None


def _sort_num(value: Any) -> float:
    numeric = _num(value)
    return -math.inf if numeric is None else numeric


def _fmt(value: Any) -> str:
    numeric = _num(value)
    return "n/a" if numeric is None else f"{numeric:.4f}"


def _latex(value: Any) -> str:
    return str(value).replace("\\", "\\textbackslash{}").replace("_", "\\_").replace("&", "\\&").replace("%", "\\%")


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