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

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

from cil.metrics import (  # noqa: E402
    MetricInputError,
    any_unsafe,
    branch_car,
    candidate_diversity,
    collapse_rate,
    macro_micro_summary,
    mean_nearest_distance_to_set,
    measured_support_gap,
    negative_near_at_threshold,
    normalized_causal_action_regret,
    outcome_ptr_at_k,
    pairwise_causal_dominance_ece,
    positives_closer_than_negatives,
    proxy_positive_tangent_coverage_at_k,
    proxy_support_distance,
    selector_regret_at_k,
    selected_unsafe,
    safety_label_coverage,
    outcome_safety_violation,
    unsafe_rate,
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Evaluate CIL/CTT metrics while keeping measured outcome metrics "
            "separate from distance-only proxy metrics."
        )
    )
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--out-dir", type=Path, required=True)
    parser.add_argument("--mode", choices=("measured", "proxy"), required=True)
    parser.add_argument("--k", type=int, default=16)
    parser.add_argument("--epsilon", type=float, default=0.0)
    parser.add_argument("--thresholds", default="0.20,0.40")
    parser.add_argument("--bootstrap-samples", type=int, default=1000)
    parser.add_argument("--confidence", type=float, default=0.95)
    parser.add_argument(
        "--no-markdown-report",
        action="store_true",
        help="Do not write report.md; use this when README.md is the only persistent Markdown file.",
    )
    args = parser.parse_args(argv)

    if args.k <= 0:
        parser.error("--k must be positive")
    thresholds = _parse_thresholds(args.thresholds)
    payload = json.loads(args.input.read_text())
    rows = payload.get("rows", payload) if isinstance(payload, dict) else payload
    if not isinstance(rows, list):
        parser.error("input must be a JSON list or an object with a rows list")

    metric_rows = []
    for index, row in enumerate(rows):
        if not isinstance(row, dict):
            raise MetricInputError(f"row {index} must be an object")
        metric_rows.append(
            _measured_row(row, k=args.k, epsilon=args.epsilon)
            if args.mode == "measured"
            else _proxy_row(row, k=args.k, thresholds=thresholds)
        )

    metric_names = sorted(
        {
            key
            for row in metric_rows
            for key, value in row.items()
            if key not in {"task_id", "seed", "chart_id", "mode"}
            and isinstance(value, (int, float))
            and math.isfinite(float(value))
        }
    )
    summary = {
        name: macro_micro_summary(
            metric_rows,
            name,
            bootstrap_samples=args.bootstrap_samples,
            confidence=args.confidence,
        )
        for name in metric_names
    }

    out_dir = args.out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "metrics.json").write_text(
        json.dumps(
            {
                "mode": args.mode,
                "k": args.k,
                "epsilon": args.epsilon,
                "thresholds": thresholds,
                "num_rows": len(metric_rows),
                "rows": metric_rows,
                "summary": summary,
            },
            indent=2,
            sort_keys=True,
        )
        + "\n"
    )
    (out_dir / "metrics_by_task.json").write_text(
        json.dumps(_group_means(metric_rows, "task_id", metric_names), indent=2, sort_keys=True)
        + "\n"
    )
    (out_dir / "metrics_by_seed.json").write_text(
        json.dumps(_group_means(metric_rows, "seed", metric_names), indent=2, sort_keys=True)
        + "\n"
    )
    (out_dir / "table.tex").write_text(_latex_table(summary) + "\n")
    _write_run_metadata(out_dir, args, payload, metric_names)
    report_path = out_dir / "report.md"
    if args.no_markdown_report:
        report_path.unlink(missing_ok=True)
    else:
        report_path.write_text(_markdown_report(args.mode, args.k, summary) + "\n")
    print(json.dumps({"out_dir": str(out_dir), "num_rows": len(metric_rows)}, indent=2))
    return 0


def _measured_row(row: dict[str, Any], *, k: int, epsilon: float) -> dict[str, Any]:
    if not bool(row.get("candidates_evaluated", False)):
        raise MetricInputError(
            "measured mode requires candidates_evaluated=true for every row; "
            "distance-only rows must use --mode proxy"
        )
    utilities = _numbers(row, "generated_utilities")
    if not utilities:
        raise MetricInputError("measured rows require generated_utilities")
    selected_index = int(row.get("selected_index", 0))
    hidden = _numbers(row, "hidden_chart_utilities", required=False)
    candidate_success = _bool_numbers(row, "candidate_success", required=False)
    base_success = _optional_bool(row.get("base_success"))
    candidate_outcomes = _outcomes(row, "candidate_outcomes", required=False)
    selected_utility = utilities[selected_index]
    prefix = utilities[:k]
    output = _base_row(row, mode="measured")
    base_utility = _number(row, "base_utility")
    proposal_oracle_utility = max(prefix)
    output[f"outcome_ptr_at_{k}"] = outcome_ptr_at_k(
        utilities,
        base_utility,
        epsilon=epsilon,
        k=k,
        candidates_evaluated=True,
    )
    output[f"selector_regret_at_{k}"] = selector_regret_at_k(
        utilities,
        selected_index=selected_index,
        k=k,
        candidates_evaluated=True,
    )
    output[f"branch_car_at_{k}"] = branch_car(max(prefix), selected_utility)
    ncar_to_proposal = _stable_ncar(
        proposal_oracle_utility,
        selected_utility,
        base_utility,
    )
    if ncar_to_proposal is not None:
        output[f"ncar_to_proposal_oracle_at_{k}"] = ncar_to_proposal
    output["base_utility"] = base_utility
    output[f"selected_utility_at_{k}"] = selected_utility
    output[f"proposal_oracle_utility_at_{k}"] = proposal_oracle_utility
    output[f"selected_utility_gain_over_base_at_{k}"] = selected_utility - base_utility
    output[f"proposal_oracle_utility_gain_over_base_at_{k}"] = (
        proposal_oracle_utility - base_utility
    )
    if base_success is not None:
        output["base_success"] = float(base_success)
    base_outcome = row.get("base_outcome")
    if isinstance(base_outcome, dict):
        base_safety = outcome_safety_violation(base_outcome)
        output["base_safety_label_known"] = float(base_safety is not None)
        if base_safety is not None:
            output["base_unsafe_known"] = float(base_safety)
    if candidate_outcomes:
        output[f"generated_safety_label_coverage_at_{k}"] = safety_label_coverage(
            candidate_outcomes,
            k=k,
        )
        generated_unsafe = unsafe_rate(candidate_outcomes, k=k)
        if generated_unsafe is not None:
            output[f"generated_unsafe_rate_known_at_{k}"] = generated_unsafe
        any_generated_unsafe = any_unsafe(candidate_outcomes, k=k)
        if any_generated_unsafe is not None:
            output[f"any_generated_unsafe_known_at_{k}"] = any_generated_unsafe
        if selected_index < min(k, len(candidate_outcomes)):
            selected_safety = outcome_safety_violation(candidate_outcomes[selected_index])
            output[f"selected_safety_label_known_at_{k}"] = float(
                selected_safety is not None
            )
            selected_safety_value = selected_unsafe(
                candidate_outcomes,
                selected_index=selected_index,
                k=k,
            )
            if selected_safety_value is not None:
                output[f"selected_unsafe_known_at_{k}"] = selected_safety_value
        if prefix:
            oracle_index = max(range(len(prefix)), key=lambda item: prefix[item])
            if oracle_index < len(candidate_outcomes):
                oracle_safety = outcome_safety_violation(candidate_outcomes[oracle_index])
                output[f"proposal_oracle_safety_label_known_at_{k}"] = float(
                    oracle_safety is not None
                )
                if oracle_safety is not None:
                    output[f"proposal_oracle_unsafe_known_at_{k}"] = float(
                        oracle_safety
                    )
    if candidate_success:
        success_prefix = candidate_success[:k]
        selected_success = float(success_prefix[selected_index])
        proposal_oracle_success = float(any(success_prefix))
        output[f"selected_success_at_{k}"] = selected_success
        output[f"proposal_oracle_success_at_{k}"] = proposal_oracle_success
        if base_success is not None:
            output[f"selected_success_gain_over_base_at_{k}"] = (
                selected_success - float(base_success)
            )
            output[f"proposal_oracle_success_gain_over_base_at_{k}"] = (
                proposal_oracle_success - float(base_success)
            )
    if hidden:
        hidden_oracle_utility = max(hidden)
        output[f"support_gap_at_{k}"] = measured_support_gap(
            hidden_oracle_utility,
            max(prefix),
            candidates_evaluated=True,
        )
        output[f"hidden_chart_oracle_utility_at_{k}"] = hidden_oracle_utility
        output[f"total_car_to_hidden_at_{k}"] = branch_car(
            hidden_oracle_utility,
            selected_utility,
        )
        ncar_to_hidden = _stable_ncar(
            hidden_oracle_utility,
            selected_utility,
            base_utility,
        )
        if ncar_to_hidden is not None:
            output[f"ncar_to_hidden_chart_oracle_at_{k}"] = ncar_to_hidden
        hidden_gap = abs(hidden_oracle_utility - base_utility)
        if hidden_gap > 0.0:
            output[f"support_gap_fraction_to_hidden_at_{k}"] = (
                output[f"support_gap_at_{k}"] / hidden_gap
            )
            output[f"selector_gap_fraction_to_hidden_at_{k}"] = (
                output[f"selector_regret_at_{k}"] / hidden_gap
            )
        if candidate_success:
            hidden_oracle_success = float(any(value >= 1.0 for value in hidden))
            output[f"hidden_chart_oracle_success_at_{k}"] = hidden_oracle_success
            output[f"success_support_gap_at_{k}"] = max(
                0.0,
                hidden_oracle_success - output[f"proposal_oracle_success_at_{k}"],
            )
            output[f"success_selector_gap_at_{k}"] = max(
                0.0,
                output[f"proposal_oracle_success_at_{k}"]
                - output[f"selected_success_at_{k}"],
            )
            output[f"success_total_car_to_hidden_at_{k}"] = max(
                0.0,
                hidden_oracle_success - output[f"selected_success_at_{k}"],
            )
    predicted = _numbers(row, "predicted_scores", required=False)
    if predicted and len(predicted) >= len(utilities):
        ece = pairwise_causal_dominance_ece(predicted[: len(utilities)], utilities)
        output["pairwise_causal_calibration_ece"] = ece["ece"]
    return output


def _proxy_row(row: dict[str, Any], *, k: int, thresholds: list[float]) -> dict[str, Any]:
    generated = _matrix(row, "generated_tangents")
    positives = _matrix(row, "positive_tangents")
    negatives = _matrix(row, "negative_tangents", required=False)
    output = _base_row(row, mode="proxy")
    for threshold in thresholds:
        suffix = _threshold_suffix(threshold)
        output[f"pptc_at_{k}_thr_{suffix}"] = proxy_positive_tangent_coverage_at_k(
            generated,
            positives,
            threshold=threshold,
            k=k,
        )
        output[f"negative_near_at_{k}_thr_{suffix}"] = negative_near_at_threshold(
            generated,
            negatives,
            threshold=threshold,
            k=k,
        )
    distance = proxy_support_distance(generated, positives, k=k)
    if distance is not None:
        output[f"proxy_support_distance_at_{k}"] = distance
    positive_distance = mean_nearest_distance_to_set(generated, positives, k=k)
    if positive_distance is not None:
        output[f"mean_positive_distance_at_{k}"] = positive_distance
    negative_distance = mean_nearest_distance_to_set(generated, negatives, k=k)
    if negative_distance is not None:
        output[f"mean_negative_distance_at_{k}"] = negative_distance
    closer = positives_closer_than_negatives(generated, positives, negatives, k=k)
    if closer is not None:
        output[f"pos_closer_than_neg_at_{k}"] = closer
    output[f"candidate_diversity_at_{k}"] = candidate_diversity(generated, k=k)
    output[f"collapse_rate_at_{k}"] = collapse_rate(generated, k=k)
    return output


def _base_row(row: dict[str, Any], *, mode: str) -> dict[str, Any]:
    return {
        "mode": mode,
        "chart_id": str(row.get("chart_id", row.get("group_id", "unknown"))),
        "task_id": str(row.get("task_id", "unknown")),
        "seed": str(row.get("seed", "unknown")),
    }


def _numbers(row: dict[str, Any], key: str, *, required: bool = True) -> list[float]:
    values = row.get(key)
    if values is None:
        if required:
            raise MetricInputError(f"row requires {key}")
        return []
    if not isinstance(values, list):
        raise MetricInputError(f"{key} must be a list")
    return [float(value) for value in values]


def _number(row: dict[str, Any], key: str) -> float:
    if key not in row:
        raise MetricInputError(f"row requires {key}")
    return float(row[key])


def _bool_numbers(row: dict[str, Any], key: str, *, required: bool = True) -> list[bool]:
    values = row.get(key)
    if values is None:
        if required:
            raise MetricInputError(f"row requires {key}")
        return []
    if not isinstance(values, list):
        raise MetricInputError(f"{key} must be a list")
    return [bool(value) for value in values]


def _optional_bool(value: Any) -> bool | None:
    if value is None:
        return None
    return bool(value)


def _stable_ncar(
    oracle_utility: float,
    selected_utility: float,
    base_utility: float,
    *,
    min_denominator: float = 1.0e-3,
) -> float | None:
    """Return NCAR only when the base-to-oracle gap is numerically meaningful."""

    if abs(float(oracle_utility) - float(base_utility)) <= min_denominator:
        return None
    return normalized_causal_action_regret(
        oracle_utility,
        selected_utility,
        base_utility,
    )


def _matrix(row: dict[str, Any], key: str, *, required: bool = True) -> list[list[float]]:
    values = row.get(key)
    if values is None:
        if required:
            raise MetricInputError(f"row requires {key}")
        return []
    if not isinstance(values, list):
        raise MetricInputError(f"{key} must be a list of vectors")
    return [[float(item) for item in vector] for vector in values]


def _outcomes(row: dict[str, Any], key: str, *, required: bool = True) -> list[dict[str, Any]]:
    values = row.get(key)
    if values is None:
        if required:
            raise MetricInputError(f"row requires {key}")
        return []
    if not isinstance(values, list):
        raise MetricInputError(f"{key} must be a list of outcome objects")
    outcomes: list[dict[str, Any]] = []
    for index, value in enumerate(values):
        if not isinstance(value, dict):
            raise MetricInputError(f"{key}[{index}] must be an outcome object")
        outcomes.append(value)
    return outcomes


def _parse_thresholds(raw: str) -> list[float]:
    values = [float(item.strip()) for item in raw.split(",") if item.strip()]
    if not values or any(value < 0.0 for value in values):
        raise ValueError("--thresholds must contain non-negative values")
    return values


def _threshold_suffix(value: float) -> str:
    return f"{value:.2f}".replace(".", "p")


def _group_means(
    rows: list[dict[str, Any]],
    key: str,
    metric_names: list[str],
) -> dict[str, dict[str, float]]:
    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        grouped[str(row.get(key, "unknown"))].append(row)
    output: dict[str, dict[str, float]] = {}
    for group, group_rows in sorted(grouped.items()):
        payload: dict[str, float] = {}
        for metric in metric_names:
            values = [
                float(row[metric])
                for row in group_rows
                if isinstance(row.get(metric), (int, float))
                and math.isfinite(float(row[metric]))
            ]
            if values:
                payload[metric] = sum(values) / len(values)
        output[group] = payload
    return output


def _write_run_metadata(
    out_dir: Path,
    args: argparse.Namespace,
    input_payload: Any,
    metric_names: list[str],
) -> None:
    data_hash = _payload_hash(input_payload)
    split_hash = _extract_hash(
        input_payload,
        (
            "split_hash",
            "target_split_hash",
            "eval_target_split_hash",
            "selector_split_hash",
        ),
    )
    if split_hash is None:
        split_hash = data_hash

    (out_dir / "config.yaml").write_text(
        "\n".join(
            [
                f"input: {args.input}",
                f"mode: {args.mode}",
                f"k: {args.k}",
                f"epsilon: {args.epsilon}",
                f"thresholds: {args.thresholds}",
                f"bootstrap_samples: {args.bootstrap_samples}",
                f"confidence: {args.confidence}",
                f"no_markdown_report: {bool(args.no_markdown_report)}",
                "metric_names:",
                *[f"  - {name}" for name in metric_names],
            ]
        )
        + "\n"
    )
    (out_dir / "command.txt").write_text(
        "python scripts/eval_metrics.py " + " ".join(sys.argv[1:]) + "\n"
    )
    (out_dir / "git_hash.txt").write_text(_git_hash() + "\n")
    (out_dir / "data_hash.txt").write_text(data_hash + "\n")
    (out_dir / "split_hash.txt").write_text(split_hash + "\n")
    (out_dir / "train.log").write_text("metric evaluation artifact; no training\n")
    (out_dir / "eval.log").write_text(
        "\n".join(
            [
                f"input={args.input}",
                f"mode={args.mode}",
                f"k={args.k}",
                f"num_metrics={len(metric_names)}",
                f"markdown_report_written={not bool(args.no_markdown_report)}",
            ]
        )
        + "\n"
    )


def _payload_hash(payload: Any) -> str:
    blob = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str).encode()
    return hashlib.sha256(blob).hexdigest()


def _extract_hash(payload: Any, keys: tuple[str, ...]) -> str | None:
    if isinstance(payload, dict):
        for key in keys:
            value = payload.get(key)
            if isinstance(value, str) and value.strip():
                return value.strip()
        for value in payload.values():
            nested = _extract_hash(value, keys)
            if nested is not None:
                return nested
    elif isinstance(payload, list):
        for value in payload:
            nested = _extract_hash(value, keys)
            if nested is not None:
                return nested
    return None


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 _latex_table(summary: dict[str, Any]) -> str:
    lines = [
        "% Auto-generated by scripts/eval_metrics.py",
        "\\begin{tabular}{lrrrr}",
        "\\toprule",
        "Metric & N & Micro mean & CI low & CI high \\\\",
        "\\midrule",
    ]
    for name, payload in sorted(summary.items()):
        micro = payload["micro"]
        lines.append(
            f"{_latex_escape(name)} & {micro['n']} & {_fmt(micro['mean'])} & "
            f"{_fmt(micro['low'])} & {_fmt(micro['high'])} \\\\"
        )
    lines.extend(["\\bottomrule", "\\end{tabular}"])
    return "\n".join(lines)


def _markdown_report(mode: str, k: int, summary: dict[str, Any]) -> str:
    lines = [
        f"# Metric Evaluation ({mode})",
        "",
        f"K: `{k}`",
        "",
        "| Metric | N | Micro mean | 95% CI | Task macro | Seed macro |",
        "| --- | ---: | ---: | ---: | ---: | ---: |",
    ]
    for name, payload in sorted(summary.items()):
        micro = payload["micro"]
        task_mean = payload["macro_by_task"]["mean"]
        seed_mean = payload["macro_by_seed"]["mean"]
        lines.append(
            f"| {name} | {micro['n']} | {_fmt(micro['mean'])} | "
            f"[{_fmt(micro['low'])}, {_fmt(micro['high'])}] | "
            f"{_fmt(task_mean)} | {_fmt(seed_mean)} |"
        )
    return "\n".join(lines)


def _rms_l2(left: list[float], right: list[float]) -> float:
    if len(left) != len(right):
        raise MetricInputError("vectors must have matching dimensions")
    if not left:
        return 0.0
    return math.sqrt(sum((a - b) ** 2 for a, b in zip(left, right, strict=True)) / len(left))


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


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


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