#!/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))

import torch  # noqa: E402

from cil.metrics import (  # noqa: E402
    candidate_diversity,
    collapse_rate,
    macro_micro_summary,
    mean_nearest_distance_to_set,
    negative_near_at_threshold,
    positives_closer_than_negatives,
    proxy_positive_tangent_coverage_at_k,
    proxy_support_distance,
)
from scripts.train_ctt import Chart, load_charts  # noqa: E402


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Evaluate train-only positive tangent memory baselines on evaluator "
            "chart positives. This is PPTC/proxy support geometry, not OutcomePTR."
        )
    )
    parser.add_argument("--source-index", type=Path, default=Path("data/cil_charts/train/index.json"))
    parser.add_argument("--target-index", type=Path, default=Path("data/cil_charts/val/index.json"))
    parser.add_argument("--out-dir", type=Path, default=Path("runs/local_atlas_val_proxy"))
    parser.add_argument(
        "--mode",
        choices=("local_atlas", "task_memory", "global_memory"),
        default="local_atlas",
    )
    parser.add_argument("--k", type=int, default=16)
    parser.add_argument("--neighbors", type=int, default=8)
    parser.add_argument("--max-target-charts", type=int, default=100000)
    parser.add_argument("--thresholds", default="0.20,0.40")
    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)

    thresholds = [float(item) for item in args.thresholds.split(",") if item.strip()]
    if args.k <= 0:
        parser.error("--k must be positive")
    if args.neighbors <= 0:
        parser.error("--neighbors must be positive")
    if any(threshold < 0.0 for threshold in thresholds):
        parser.error("--thresholds must be non-negative")

    source_charts, source_index = load_charts(args.source_index, max_charts=None)
    target_charts, target_index = load_charts(args.target_index, max_charts=args.max_target_charts)
    _validate_indexes(args.source_index, source_index, args.target_index, target_index)
    source_by_task: dict[str, list[Chart]] = {}
    for chart in source_charts:
        source_by_task.setdefault(chart.task_id, []).append(chart)

    rows = []
    for target in target_charts:
        proposals = _propose(
            target,
            source_charts=source_charts,
            source_by_task=source_by_task,
            mode=args.mode,
            k=args.k,
            neighbors=args.neighbors,
        )
        rows.append(
            _metric_row(
                target=target,
                proposals=[proposal.cpu().tolist() for proposal in proposals],
                thresholds=thresholds,
                k=args.k,
            )
        )

    metric_names = sorted(
        {
            key
            for row in rows
            for key, value in row.items()
            if isinstance(value, (int, float)) and math.isfinite(float(value))
        }
        - {"num_proposals"}
    )
    summary = {name: macro_micro_summary(rows, name, bootstrap_samples=500) for name in metric_names}
    out_dir = args.out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    _write_run_provenance(out_dir, args, source_index, target_index)
    metrics = {
        "report_type": "positive_memory_proxy_eval",
        "method": args.mode,
        "k": args.k,
        "thresholds": thresholds,
        "num_rows": len(rows),
        "rows": rows,
        "summary": summary,
        "data_hash": source_index.get("content_hash"),
        "split_hash": target_index.get("split_hash"),
        "target_data_hash": target_index.get("content_hash"),
        "target_split_hash": target_index.get("split_hash"),
    }
    (out_dir / "metrics.json").write_text(json.dumps(metrics, indent=2, sort_keys=True) + "\n")
    (out_dir / "metrics_by_task.json").write_text(
        json.dumps(_by_group(rows, metric_names, "task_id"), indent=2, sort_keys=True) + "\n"
    )
    (out_dir / "metrics_by_seed.json").write_text(
        json.dumps(_by_group(rows, metric_names, "seed"), indent=2, sort_keys=True) + "\n"
    )
    (out_dir / "train.log").write_text("train-only measured positive tangent memory; no learned training\n")
    (out_dir / "eval.log").write_text(
        "\n".join(
            [
                f"source_charts={len(source_charts)} target_charts={len(target_charts)} k={args.k}",
                f"mode={args.mode} neighbors={args.neighbors}",
                f"source_index={args.source_index}",
                f"target_index={args.target_index}",
            ]
        )
        + "\n"
    )
    (out_dir / "table.tex").write_text(_table(summary) + "\n")
    _write_markdown_report(
        out_dir,
        args.mode,
        args.k,
        summary,
        no_markdown_report=args.no_markdown_report,
    )
    print(json.dumps({"out_dir": str(out_dir), "num_rows": len(rows)}, indent=2))
    return 0


def _propose(
    target: Chart,
    *,
    source_charts: list[Chart],
    source_by_task: dict[str, list[Chart]],
    mode: str,
    k: int,
    neighbors: int,
) -> list[torch.Tensor]:
    if mode == "global_memory":
        pool = source_charts
    else:
        pool = source_by_task.get(target.task_id, source_charts)
    if mode == "task_memory":
        ranked = sorted(pool, key=lambda chart: chart.chart_id)
    else:
        ranked = sorted(
            pool,
            key=lambda chart: torch.linalg.vector_norm(chart.feature - target.feature).item(),
        )
        if mode == "local_atlas":
            ranked = ranked[:neighbors]

    proposals: list[torch.Tensor] = []
    for chart in ranked:
        for positive in chart.positives:
            proposals.append(positive)
            if len(proposals) >= k:
                return proposals
    return proposals


def _metric_row(
    *,
    target: Chart,
    proposals: list[list[float]],
    thresholds: list[float],
    k: int,
) -> dict[str, Any]:
    positives = target.positives.cpu().tolist()
    negatives = target.negatives.cpu().tolist()
    row: dict[str, Any] = {
        "chart_id": target.chart_id,
        "task_id": target.task_id,
        "seed": target.seed,
        "num_proposals": len(proposals),
    }
    for threshold in thresholds:
        suffix = f"{threshold:.2f}".replace(".", "p")
        row[f"pptc_at_{k}_thr_{suffix}"] = proxy_positive_tangent_coverage_at_k(
            proposals,
            positives,
            threshold=threshold,
            k=k,
        )
        row[f"negative_near_at_{k}_thr_{suffix}"] = negative_near_at_threshold(
            proposals,
            negatives,
            threshold=threshold,
            k=k,
        )
    row[f"proxy_support_distance_at_{k}"] = proxy_support_distance(proposals, positives, k=k)
    row[f"mean_positive_distance_at_{k}"] = mean_nearest_distance_to_set(proposals, positives, k=k)
    row[f"mean_negative_distance_at_{k}"] = mean_nearest_distance_to_set(proposals, negatives, k=k)
    row[f"pos_closer_than_neg_at_{k}"] = positives_closer_than_negatives(
        proposals,
        positives,
        negatives,
        k=k,
    )
    row[f"candidate_diversity_at_{k}"] = candidate_diversity(proposals, k=k)
    row[f"collapse_rate_at_{k}"] = collapse_rate(proposals, k=k)
    return row


def _validate_indexes(
    source_path: Path,
    source_index: dict[str, Any],
    target_path: Path,
    target_index: dict[str, Any],
) -> None:
    if source_index.get("split") != "train" or not source_index.get("retrieval_index_allowed"):
        raise SystemExit(f"{source_path} must be the train-only retrieval index")
    if not source_index.get("include_outcomes"):
        raise SystemExit(f"{source_path} must include train outcomes for positive memory")
    if not target_index.get("include_outcomes"):
        raise SystemExit(f"{target_path} must include evaluator outcomes for PPTC labels")
    if target_index.get("split") != "train" and target_index.get("retrieval_index_allowed"):
        raise SystemExit(f"{target_path} is non-train but marked retrieval_index_allowed")


def _write_run_provenance(
    out_dir: Path,
    args: argparse.Namespace,
    source_index: dict[str, Any],
    target_index: dict[str, Any],
) -> None:
    (out_dir / "config.yaml").write_text("\n".join(f"{k}: {v}" for k, v in sorted(vars(args).items())) + "\n")
    (out_dir / "command.txt").write_text(
        "python scripts/eval_chart_positive_memory_proxy.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(source_index.get("content_hash", "")) + "\n")
    (out_dir / "split_hash.txt").write_text(str(target_index.get("split_hash", "")) + "\n")


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


def _by_group(
    rows: list[dict[str, Any]],
    metric_names: list[str],
    group_key: str,
) -> dict[str, dict[str, float]]:
    output: dict[str, dict[str, float]] = {}
    for row in rows:
        group = str(row[group_key])
        output.setdefault(group, {})
    for group in output:
        group_rows = [row for row in rows if str(row[group_key]) == group]
        for metric in metric_names:
            values = [float(row[metric]) for row in group_rows if isinstance(row.get(metric), (int, float))]
            if values:
                output[group][metric] = sum(values) / len(values)
    return output


def _table(summary: dict[str, Any]) -> str:
    lines = [
        "% Auto-generated by scripts/eval_chart_positive_memory_proxy.py",
        "\\begin{tabular}{lrrr}",
        "\\toprule",
        "Metric & N & Mean & CI high \\\\",
        "\\midrule",
    ]
    for name, payload in sorted(summary.items()):
        micro = payload["micro"]
        lines.append(
            f"{_latex_escape(name)} & {micro['n']} & {micro['mean']:.4f} & "
            f"{micro['high']:.4f} \\\\"
        )
    lines.extend(["\\bottomrule", "\\end{tabular}"])
    return "\n".join(lines)


def _report(method: str, k: int, summary: dict[str, Any]) -> str:
    lines = [
        "# Positive Memory Proxy Evaluation",
        "",
        f"Method: `{method}`",
        f"K: `{k}`",
        "",
        "| Metric | N | Mean | 95% CI |",
        "| --- | ---: | ---: | ---: |",
    ]
    for name, payload in sorted(summary.items()):
        micro = payload["micro"]
        lines.append(
            f"| {name} | {micro['n']} | {micro['mean']:.4f} | "
            f"[{micro['low']:.4f}, {micro['high']:.4f}] |"
        )
    lines.append("")
    lines.append("This is PPTC/proxy support geometry, not OutcomePTR or rollout success.")
    return "\n".join(lines)


def _write_markdown_report(
    out_dir: Path,
    method: str,
    k: int,
    summary: 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(method, k, summary) + "\n")


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


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