#!/usr/bin/env python3
"""Contributor-held-out evaluation harness for the Timmy Bristol classifier.

This is the comparison and release-gate tool for issue #29: it turns candidate predictions into the
numbers the ROADMAP gates are written in — macro-F1, per-class precision/recall, a confusion matrix,
calibration error, coverage, and precision-at-coverage — and it refuses a corpus whose splits are
not contributor-held-out or not reproducible.

Nothing here trains, scores, or diagnoses: it reads predictions that something else produced, and it
reports what those predictions would have done on held-out contributors. Every fixture used to test
it is synthetic.
"""
from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path

BRISTOL_CLASSES = tuple(range(1, 8))
SPLITS = ("train", "validation", "test")
# The split rule is the ingestion rule (scripts/ingest_training_photo.py): the bucket of the first
# byte of the pseudonymous subject key. Keeping it here means a split can be re-derived from the
# key alone, so a corpus cannot quietly re-label its own splits.
_TRAIN_EDGE = 205
_VALIDATION_EDGE = 230
_SHA_LIKE = re.compile(r"^[0-9a-f]{16}$")

# The ROADMAP release gates. A candidate must clear every one on its held-out split; a limit is
# inclusive, so a model exactly on the line passes.
RELEASE_GATES = {"macroF1": 0.80, "precisionAtCoverage": 0.90, "ece": 0.05}


class EvaluationError(Exception):
    """The corpus or the caller asked for a comparison that cannot be trusted."""


def expected_split(subject_key: str) -> str:
    """The split a pseudonymous subject key must belong to, by the ingestion rule."""
    if not isinstance(subject_key, str) or not _SHA_LIKE.match(subject_key):
        raise EvaluationError("subject key must be the 16-character pseudonymous digest")
    bucket = int(subject_key[:2], 16)
    if bucket < _TRAIN_EDGE:
        return "train"
    if bucket < _VALIDATION_EDGE:
        return "validation"
    return "test"


def _calibration_error(accepted: list[tuple[int, int, float]], bins: int = 10) -> tuple[float, list[dict]]:
    """Expected calibration error over the accepted suggestions, by equal-width confidence bin."""
    if not accepted:
        return 0.0, []
    edges = [index / bins for index in range(bins + 1)]
    buckets: list[list[tuple[int, int, float]]] = [[] for _ in range(bins)]
    for truth, predicted, confidence in accepted:
        index = min(int(confidence * bins), bins - 1)
        buckets[index].append((truth, predicted, confidence))
    error = 0.0
    detail: list[dict] = []
    for index, bucket in enumerate(buckets):
        if not bucket:
            continue
        accuracy = sum(1 for truth, predicted, _ in bucket if truth == predicted) / len(bucket)
        mean_confidence = sum(confidence for _, _, confidence in bucket) / len(bucket)
        share = len(bucket) / len(accepted)
        error += share * abs(accuracy - mean_confidence)
        detail.append(
            {
                "lower": edges[index],
                "upper": edges[index + 1],
                "count": len(bucket),
                "accuracy": accuracy,
                "meanConfidence": mean_confidence,
            }
        )
    return error, detail


def assert_contributor_held_out(records: list[dict]) -> None:
    """Refuse a corpus whose splits are not contributor-held-out and reproducible.

    A split must be re-derivable from the pseudonymous subject key alone (the ingestion rule), and
    one contributor's photos must all land in the same split. Otherwise a held-out score is not
    held out: the model has already seen that person's body in training.
    """
    splits_by_subject: dict[str, set[str]] = {}
    for entry in records:
        subject = entry.get("subjectKey")
        if not isinstance(subject, str) or not subject:
            raise EvaluationError("every prediction must carry the pseudonymous subject key")
        split = entry.get("split")
        if split not in SPLITS:
            raise EvaluationError(f"split must be one of {', '.join(SPLITS)}: {split!r}")
        splits_by_subject.setdefault(subject, set()).add(split)
    for subject, splits in splits_by_subject.items():
        if len(splits) > 1:
            raise EvaluationError(
                f"contributor {subject} leaks across {', '.join(sorted(splits))}: splits must be contributor-held-out"
            )
    for entry in records:
        subject = entry["subjectKey"]
        required = expected_split(subject)
        if entry["split"] != required:
            raise EvaluationError(
                f"split {entry['split']!r} for contributor {subject} is not reproducible: the subject key derives {required!r}"
            )


def compare(models: dict[str, list[dict]], *, min_coverage: float = 0.5, gates: dict | None = None) -> dict:
    """Compare candidate models against the release gates.

    Each candidate gets one threshold, chosen on its own validation split, and is then measured on
    its test split. The gates are the ROADMAP release gates; all limits are inclusive.
    """
    limits = dict(gates or RELEASE_GATES)
    rows: list[dict] = []
    for model_id in sorted(models):
        records = models[model_id]
        assert_contributor_held_out(records)
        test = [entry for entry in records if entry.get("split") == "test"]
        validation = [entry for entry in records if entry.get("split") == "validation"]
        if not test:
            raise EvaluationError(f"{model_id} has no test-split predictions to measure a release on")
        if not validation:
            raise EvaluationError(f"{model_id} has no validation-split predictions to fix a threshold on")
        selection = select_threshold(validation, min_coverage=min_coverage)
        scored = evaluate(test, threshold=selection["threshold"])
        failed = []
        if scored["macroF1"] < limits["macroF1"]:
            failed.append("macroF1")
        if scored["precisionAtCoverage"] < limits["precisionAtCoverage"]:
            failed.append("precisionAtCoverage")
        if scored["ece"] > limits["ece"]:
            failed.append("ece")
        rows.append(
            {
                "model": model_id,
                **scored,
                "selectionSplit": "validation",
                "evaluatedSplit": "test",
                "minCoverage": min_coverage,
                "gate": {"passed": not failed, "failed": failed},
            }
        )
    ordered = sorted(
        rows,
        key=lambda row: (
            not row["gate"]["passed"],
            -row["macroF1"],
            -row["precisionAtCoverage"],
            row["model"],
        ),
    )
    passing = [row for row in ordered if row["gate"]["passed"]]
    return {
        "gates": limits,
        "models": ordered,
        "recommended": passing[0]["model"] if passing else None,
        "verdict": "GO" if passing else "BLOCKED",
    }


def select_threshold(records: list[dict], *, min_coverage: float = 0.5) -> dict:
    """Choose the abstention threshold that maximises precision at a required coverage.

    Selection reads the validation split only. Fixing a threshold on the same data the release gate
    is measured on would manufacture whatever precision the report needs.
    """
    if not records:
        raise EvaluationError("threshold selection needs validation predictions")
    if not 0.0 <= min_coverage <= 1.0:
        raise EvaluationError(f"a coverage floor must be between 0 and 1: {min_coverage}")
    if any(entry.get("split") != "validation" for entry in records):
        raise EvaluationError("threshold selection must use the validation split only")
    assert_contributor_held_out(records)

    candidates = {0.0}
    candidates.update(
        entry["prediction"]["confidence"] for entry in records if entry["prediction"] is not None
    )
    best: dict | None = None
    for threshold in sorted(candidates):
        report = evaluate(records, threshold=threshold)
        # Float comparison: a coverage exactly on the floor must not be lost to rounding.
        if report["coverage"] + 1e-12 < min_coverage:
            continue
        if best is None:
            best = report
            continue
        more_precise = round(report["precisionAtCoverage"], 12) > round(best["precisionAtCoverage"], 12)
        equally_precise_but_wider = (
            round(report["precisionAtCoverage"], 12) == round(best["precisionAtCoverage"], 12)
            and report["coverage"] > best["coverage"]
        )
        if more_precise or equally_precise_but_wider:
            best = report
    if best is None:
        raise EvaluationError(f"no threshold reaches the {min_coverage:.0%} coverage floor on validation")
    return {**best, "selectionSplit": "validation", "minCoverage": min_coverage}


def evaluate(records: list[dict], *, threshold: float = 0.0) -> dict:
    """Score one model's accepted suggestions at a fixed abstention threshold."""
    accepted: list[tuple[int, int, float]] = []
    abstentions = {"low-confidence": 0, "no-suggestion": 0}
    for entry in records:
        truth = entry["truth"]["bristolType"]
        prediction = entry["prediction"]
        if prediction is None:
            abstentions["no-suggestion"] += 1
            continue
        if prediction["confidence"] < threshold:
            abstentions["low-confidence"] += 1
            continue
        accepted.append((truth, prediction["bristolType"], prediction["confidence"]))

    confusion: dict[str, int] = {}
    for truth, predicted, _ in accepted:
        key = f"{truth}->{predicted}"
        confusion[key] = confusion.get(key, 0) + 1

    truth_classes = {truth for truth, _, _ in accepted}
    per_class: dict[int, dict] = {}
    for label in sorted(truth_classes):
        true_positive = sum(1 for truth, predicted, _ in accepted if truth == label and predicted == label)
        predicted_positive = sum(1 for _, predicted, _ in accepted if predicted == label)
        actual_positive = sum(1 for truth, _, _ in accepted if truth == label)
        precision = true_positive / predicted_positive if predicted_positive else 0.0
        recall = true_positive / actual_positive if actual_positive else 0.0
        f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0
        per_class[label] = {
            "precision": precision,
            "recall": recall,
            "f1": f1,
            "support": actual_positive,
            "predicted": predicted_positive,
        }

    macro_f1 = sum(row["f1"] for row in per_class.values()) / len(per_class) if per_class else 0.0
    ece, calibration_bins = _calibration_error(accepted)
    truth_classes_in_split = {entry["truth"]["bristolType"] for entry in records}
    return {
        "threshold": threshold,
        "counts": {"total": len(records), "accepted": len(accepted), "abstained": len(records) - len(accepted)},
        "abstentions": abstentions,
        "coverage": len(accepted) / len(records) if records else 0.0,
        "precisionAtCoverage": (
            sum(1 for truth, predicted, _ in accepted if truth == predicted) / len(accepted) if accepted else 0.0
        ),
        "perClass": per_class,
        "excludedClasses": sorted(truth_classes_in_split - set(per_class)),
        "macroF1": macro_f1,
        "ece": ece,
        "calibrationBins": calibration_bins,
        "confusion": dict(sorted(confusion.items())),
    }


def _validate_prediction(entry: object, where: str) -> dict:
    """Normalise and bound one prediction record, or refuse it."""
    if not isinstance(entry, dict):
        raise EvaluationError(f"{where} is not a prediction object")
    model = entry.get("model")
    if not isinstance(model, str) or not model.strip():
        raise EvaluationError(f"{where} needs the model that produced the prediction")
    truth = entry.get("truth")
    if not isinstance(truth, dict) or not isinstance(truth.get("bristolType"), int):
        raise EvaluationError(f"{where} needs truth.bristolType as a whole number")
    if truth["bristolType"] not in BRISTOL_CLASSES:
        raise EvaluationError(f"{where} truth.bristolType {truth['bristolType']} is outside Bristol 1-7")
    split = entry.get("split")
    if split not in SPLITS:
        raise EvaluationError(f"{where} split must be one of {', '.join(SPLITS)}")
    subject = entry.get("subjectKey")
    if not isinstance(subject, str) or not _SHA_LIKE.match(subject):
        raise EvaluationError(f"{where} needs the 16-character pseudonymous subject key")
    prediction = entry.get("prediction")
    if prediction is not None:
        if not isinstance(prediction, dict):
            raise EvaluationError(f"{where} prediction must be an object or null")
        if prediction.get("bristolType") not in BRISTOL_CLASSES:
            raise EvaluationError(f"{where} predicted bristolType must be 1-7")
        confidence = prediction.get("confidence")
        if not isinstance(confidence, (int, float)) or isinstance(confidence, bool) or not 0.0 <= confidence <= 1.0:
            raise EvaluationError(f"{where} confidence must be a number between 0 and 1")
    return {
        "model": model.strip(),
        "subjectKey": subject,
        "split": split,
        "truth": {"bristolType": truth["bristolType"]},
        "prediction": None
        if prediction is None
        else {"bristolType": prediction["bristolType"], "confidence": float(prediction["confidence"])},
    }


def load_predictions(paths: list[str]) -> list[dict]:
    """Read prediction JSONL files, refusing anything the gates could not honestly be computed from."""
    records: list[dict] = []
    for raw_path in paths:
        path = Path(raw_path)
        if not path.is_file():
            raise EvaluationError(f"predictions file not found: {path}")
        for number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
            if not line.strip():
                continue
            try:
                entry = json.loads(line)
            except json.JSONDecodeError as exc:
                raise EvaluationError(f"{path}:{number} is not valid JSON: {exc.msg}") from exc
            records.append(_validate_prediction(entry, f"{path}:{number}"))
    if not records:
        raise EvaluationError("no prediction records were supplied")
    return records


def group_by_model(records: list[dict]) -> dict[str, list[dict]]:
    models: dict[str, list[dict]] = {}
    for entry in records:
        models.setdefault(entry["model"], []).append(
            {key: value for key, value in entry.items() if key != "model"}
        )
    return models


def format_report(report: dict, *, source: str) -> str:
    """Render the comparison as a compact receipt a human can read in a terminal."""
    lines = [
        "TIMMY CALIBRATION REPORT",
        f"source: {source}",
        f"threshold: chosen on validation only · scored on test · coverage floor {report['models'][0]['minCoverage']:.0%}"
        if report["models"]
        else "threshold: nothing to score",
        "candidates: " + ", ".join(row["model"] for row in report["models"]),
        "",
        f"{'model':<26}{'thr':>5}{'macroF1':>9}{'prec@cov':>10}{'coverage':>10}{'ECE':>7}  gate",
    ]
    for row in report["models"]:
        gate = "PASS" if row["gate"]["passed"] else "FAIL: " + ", ".join(row["gate"]["failed"])
        lines.append(
            f"{row['model']:<26}{row['threshold']:>5.2f}{row['macroF1']:>9.3f}"
            f"{row['precisionAtCoverage']:>10.3f}{row['coverage']:>10.2f}{row['ece']:>7.3f}  {gate}"
        )
    limits = report["gates"]
    lines += [
        "",
        "verdict: " + report["verdict"] + (f" — recommended {report['recommended']}" if report["recommended"] else ""),
        f"gates: macroF1 >= {limits['macroF1']:.2f} · precision@coverage >= {limits['precisionAtCoverage']:.2f} · ECE <= {limits['ece']:.2f}",
    ]
    return "\n".join(lines)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Score Bristol candidates against the release gates.")
    parser.add_argument("--predictions", action="append", required=True, metavar="PATH", help="prediction JSONL (repeatable)")
    parser.add_argument("--min-coverage", type=float, default=0.5, help="coverage floor for threshold selection")
    parser.add_argument("--json", metavar="PATH", help="write the full report as JSON here")
    parser.add_argument("--source", default=None, help="what the predictions are, recorded in the receipt")
    args = parser.parse_args(argv)

    try:
        records = load_predictions(args.predictions)
        report = compare(group_by_model(records), min_coverage=args.min_coverage)
    except EvaluationError as exc:
        print(f"calibration harness refused: {exc}", file=sys.stderr)
        return 2

    source = args.source or f"{len(records)} prediction records from {len(args.predictions)} file(s)"
    print(format_report(report, source=source))
    if args.json:
        Path(args.json).write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    return 0 if report["verdict"] == "GO" else 3


if __name__ == "__main__":
    sys.exit(main())