#!/usr/bin/env python3
"""Build a locked development report from schema-v2 token predictions."""

from __future__ import annotations

import argparse
from collections import defaultdict
from dataclasses import replace
import json
from pathlib import Path
import sys

import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from baselines import TokenPredictionRow  # noqa: E402
from evaluation.provenance import (  # noqa: E402
    git_provenance,
    validate_prediction_provenance,
)
from evaluation.metrics import binary_metrics  # noqa: E402
from evaluation.statistics import (  # noqa: E402
    assert_common_support,
    descriptive_metrics,
    observations_from_predictions,
    paired_cluster_bootstrap,
)


def parse_run_spec(value: str) -> tuple[str, int, Path]:
    try:
        method, seed_text, path_text = value.split(":", 2)
        return method, int(seed_text), Path(path_text)
    except Exception as exc:
        raise argparse.ArgumentTypeError(
            "run must be METHOD:SEED:PREDICTIONS_JSONL"
        ) from exc


def parse_args(argv=None):
    parser = argparse.ArgumentParser()
    parser.add_argument("--run", action="append", type=parse_run_spec, required=True)
    parser.add_argument(
        "--canonical-metadata-run",
        type=parse_run_spec,
        required=True,
        help=(
            "METHOD:SEED:PREDICTIONS_JSONL run whose complete source and "
            "split-group metadata is applied to the common token support"
        ),
    )
    parser.add_argument("--primary-method", required=True)
    parser.add_argument("--comparator-method", required=True)
    parser.add_argument("--bootstrap-draws", type=int, default=10_000)
    parser.add_argument("--bootstrap-seed", type=int, default=20260806)
    parser.add_argument("--expected-seeds", nargs="+", type=int, required=True)
    parser.add_argument("--expected-model-revision", required=True)
    parser.add_argument("--expected-input-sha256", required=True)
    parser.add_argument("--expected-split-sha256", required=True)
    parser.add_argument("--expected-git-commit", required=True)
    parser.add_argument("--expected-operating-fpr", type=float, required=True)
    parser.add_argument(
        "--expected-primary-protocol", default="latticespike_causal_token_v2"
    )
    parser.add_argument(
        "--expected-comparator-protocol", default="causal_token_v1"
    )
    parser.add_argument("--output-json", type=Path, required=True)
    parser.add_argument("--output-md", type=Path, required=True)
    return parser.parse_args(argv)


def load_prediction_rows(path: Path, method: str) -> list[TokenPredictionRow]:
    rows = []
    with path.open(encoding="utf-8") as file:
        for line_number, line in enumerate(file, start=1):
            if not line.strip():
                continue
            try:
                row = TokenPredictionRow.from_json_dict(json.loads(line))
            except Exception as exc:
                raise ValueError(f"{path}:{line_number} is not a prediction row") from exc
            if row.method != method:
                continue
            split = str(row.metadata.get("split", "")).lower()
            if split and split not in {"eval", "test"}:
                continue
            rows.append(row)
    if not rows:
        raise ValueError(f"{path} contains no held-out rows for method {method}")
    validate_prediction_provenance(rows)
    return rows


def operating_point_metrics(
    predictions: list[TokenPredictionRow], observations
) -> dict:
    thresholds = {
        float(row.metadata["operating_point_threshold"])
        for row in predictions
        if row.metadata.get("operating_point_threshold") is not None
    }
    if len(thresholds) != 1:
        raise ValueError(
            "each run must record one validation-locked operating-point threshold"
        )
    threshold = next(iter(thresholds))
    result = binary_metrics(
        [observation.score for observation in observations],
        [observation.label for observation in observations],
        threshold=threshold,
    )
    return result.to_json_dict()


def validate_run_contract(args, method: str, seed: int, predictions) -> None:
    expected_protocol = (
        args.expected_primary_protocol
        if method == args.primary_method
        else args.expected_comparator_protocol
    )
    errors = set()
    for row in predictions:
        metadata = row.metadata
        protocol = metadata.get("evaluation_protocol") or {}
        if protocol.get("suite") != "causal_matched":
            errors.add("suite")
        if protocol.get("protocol_id") != expected_protocol:
            errors.add("protocol_id")
        if int(metadata.get("training_seed", -1)) != int(seed):
            errors.add("training_seed")
        if metadata.get("git_commit") != args.expected_git_commit:
            errors.add("git_commit")
        if metadata.get("input_sha256") != args.expected_input_sha256:
            errors.add("input_sha256")
        if metadata.get("split_sha256") != args.expected_split_sha256:
            errors.add("split_sha256")
        if (
            metadata.get("resolved_scoring_model_revision")
            != args.expected_model_revision
        ):
            errors.add("model_revision")
        if metadata.get("threshold_selection_split") != "validation":
            errors.add("threshold_selection_split")
        if metadata.get("threshold_selection_policy") != "validation_target_fpr":
            errors.add("threshold_selection_policy")
        target = metadata.get("target_false_positive_rate")
        if target is None or not np.isclose(
            float(target), float(args.expected_operating_fpr)
        ):
            errors.add("target_false_positive_rate")
    if errors:
        raise ValueError(
            f"run contract mismatch for {method} seed {seed}: {sorted(errors)}"
        )


def bind_canonical_metadata(observations, canonical_observations):
    """Bind method-independent source/group fields by exact token support."""
    assert_common_support(
        {"candidate": observations, "canonical": canonical_observations}
    )
    metadata_by_support = {}
    for observation in canonical_observations:
        metadata = (observation.source, observation.split_group)
        previous = metadata_by_support.setdefault(observation.support_key, metadata)
        if previous != metadata:
            raise ValueError("canonical metadata conflicts within token support")
        if observation.source == "unknown":
            raise ValueError("canonical metadata contains unknown source")
        if observation.split_group.startswith("row:"):
            raise ValueError("canonical metadata is missing dataset split_group")
    return [
        replace(
            observation,
            source=metadata_by_support[observation.support_key][0],
            split_group=metadata_by_support[observation.support_key][1],
        )
        for observation in observations
    ]


def build_report(args) -> dict:
    canonical_spec = args.canonical_metadata_run
    if canonical_spec not in args.run:
        raise ValueError("--canonical-metadata-run must also appear in --run")
    canonical_method, canonical_seed, canonical_path = canonical_spec
    canonical_predictions = load_prediction_rows(canonical_path, canonical_method)
    validate_run_contract(
        args, canonical_method, canonical_seed, canonical_predictions
    )
    canonical_observations = observations_from_predictions(canonical_predictions)
    bind_canonical_metadata(canonical_observations, canonical_observations)

    method_runs = defaultdict(list)
    run_rows = []
    for method, seed, path in args.run:
        predictions = load_prediction_rows(path, method)
        validate_run_contract(args, method, seed, predictions)
        observations = bind_canonical_metadata(
            observations_from_predictions(predictions), canonical_observations
        )
        metrics = descriptive_metrics(observations)
        metrics["operating_point"] = operating_point_metrics(
            predictions, observations
        )
        method_runs[method].append(observations)
        run_rows.append(
            {
                "method": method,
                "seed": seed,
                "predictions": str(path),
                "provenance": validate_prediction_provenance(predictions),
                "metrics": metrics,
            }
        )
    if args.primary_method not in method_runs or args.comparator_method not in method_runs:
        raise ValueError("primary and comparator methods must both have run specs")
    seen_specs = [(method, seed) for method, seed, _ in args.run]
    if len(seen_specs) != len(set(seen_specs)):
        raise ValueError("duplicate method/seed run specs")
    for method in (args.primary_method, args.comparator_method):
        observed = sorted(seed for candidate, seed in seen_specs if candidate == method)
        if observed != sorted(args.expected_seeds):
            raise ValueError(
                f"{method} seeds {observed} do not match {sorted(args.expected_seeds)}"
            )
    aggregate = {}
    for method in sorted(method_runs):
        rows = [row for row in run_rows if row["method"] == method]
        metric_names = (
            "micro_average_precision",
            "micro_auroc",
            "answer_balanced_average_precision",
            "source_macro_average_precision",
        )
        metric_summary = {}
        for metric_name in metric_names:
            values = [float(row["metrics"][metric_name]) for row in rows]
            metric_summary[f"{metric_name}_mean"] = float(np.mean(values))
            metric_summary[f"{metric_name}_sd"] = (
                float(np.std(values, ddof=1)) if len(values) > 1 else 0.0
            )
        for metric_name in (
            "threshold",
            "precision",
            "recall",
            "f1",
            "false_positive_rate",
        ):
            values = [
                float(row["metrics"]["operating_point"][metric_name])
                for row in rows
            ]
            metric_summary[f"operating_{metric_name}_mean"] = float(
                np.mean(values)
            )
            metric_summary[f"operating_{metric_name}_sd"] = (
                float(np.std(values, ddof=1)) if len(values) > 1 else 0.0
            )
        aggregate[method] = {
            "n_seeds": len(rows),
            **metric_summary,
            "runs": rows,
        }
    bootstrap = paired_cluster_bootstrap(
        method_runs,
        primary_method=args.primary_method,
        comparator_method=args.comparator_method,
        draws=args.bootstrap_draws,
        seed=args.bootstrap_seed,
    )
    return {
        "status": "v5.3 development result; not untouched confirmatory evidence",
        "prediction_schema_version": 2,
        "report_git": git_provenance(REPO_ROOT, require_clean=True),
        "canonical_metadata": {
            "method": canonical_method,
            "seed": canonical_seed,
            "predictions": str(canonical_path),
            "sources": sorted(
                {observation.source for observation in canonical_observations}
            ),
            "split_groups": len(
                {observation.split_group for observation in canonical_observations}
            ),
            "provenance": validate_prediction_provenance(
                canonical_predictions
            ),
        },
        "primary_estimand": bootstrap["estimand"],
        "methods": aggregate,
        "paired_source_stratified_group_bootstrap": bootstrap,
    }


def write_markdown(path: Path, report: dict) -> None:
    lines = [
        "# Paper-ready v5.3 development evaluation",
        "",
        "> v5.3 is development data with judge-derived labels; this is not an untouched confirmatory result.",
        "",
        f"- Reporting commit: `{report['report_git']['commit']}`",
        f"- Canonical metadata run: `{report['canonical_metadata']['method']}` seed {report['canonical_metadata']['seed']}",
        f"- Canonical sources/groups: {len(report['canonical_metadata']['sources'])} / {report['canonical_metadata']['split_groups']}",
        "",
        "| Method | Seeds | Micro AP | Micro AUROC | Answer-balanced AP | Source-macro AP |",
        "|---|---:|---:|---:|---:|---:|",
    ]
    for method, values in report["methods"].items():
        lines.append(
            f"| {method} | {values['n_seeds']} | "
            f"{values['micro_average_precision_mean']:.6f} ± {values['micro_average_precision_sd']:.6f} | "
            f"{values['micro_auroc_mean']:.6f} ± {values['micro_auroc_sd']:.6f} | "
            f"{values['answer_balanced_average_precision_mean']:.6f} ± {values['answer_balanced_average_precision_sd']:.6f} | "
            f"{values['source_macro_average_precision_mean']:.6f} ± {values['source_macro_average_precision_sd']:.6f} |"
        )
    lines.extend(
        [
            "",
            "## Validation-locked operating point on test",
            "",
            "| Method | Threshold | Precision | Recall | F1 | Actual FPR |",
            "|---|---:|---:|---:|---:|---:|",
        ]
    )
    for method, values in report["methods"].items():
        lines.append(
            f"| {method} | "
            f"{values['operating_threshold_mean']:.6f} ± {values['operating_threshold_sd']:.6f} | "
            f"{values['operating_precision_mean']:.6f} ± {values['operating_precision_sd']:.6f} | "
            f"{values['operating_recall_mean']:.6f} ± {values['operating_recall_sd']:.6f} | "
            f"{values['operating_f1_mean']:.6f} ± {values['operating_f1_sd']:.6f} | "
            f"{values['operating_false_positive_rate_mean']:.6f} ± {values['operating_false_positive_rate_sd']:.6f} |"
        )
    bootstrap = report["paired_source_stratified_group_bootstrap"]
    low, high = bootstrap["ci_95_percentile"]
    lines.extend(
        [
            "",
            "## Primary paired comparison",
            "",
            f"- Estimand: {bootstrap['estimand']}",
            f"- Mean AP delta: {bootstrap['delta_mean']:.6f}",
            f"- 95% cluster-bootstrap CI: [{low:.6f}, {high:.6f}]",
            f"- Bootstrap draws used: {bootstrap['draws_used']:,}",
            "- Resampling unit: complete duplicate split groups, sampled within source.",
            "- Seed variation is reported separately; seed and token observations are not pooled.",
        ]
    )
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def main(argv=None):
    args = parse_args(argv)
    report = build_report(args)
    args.output_json.parent.mkdir(parents=True, exist_ok=True)
    args.output_json.write_text(
        json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    write_markdown(args.output_md, report)
    print(args.output_json)
    print(args.output_md)


if __name__ == "__main__":
    main()