#!/usr/bin/env python3
"""Re-evaluate saved segmentation predictions at multiple maxDets values.

This is a deliberately narrow companion to ``reevaluate_dense_predictions``.
It evaluates only the full image set and segmentation IoU type, avoiding the
unneeded bbox and density-stratum passes when extending a maxDets ablation.
No prediction, model weight, or annotation is modified.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
from collections import Counter
from pathlib import Path

from reevaluate_dense_predictions import load_json, run_coco_eval_multi_max, sha256


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--ann-file", type=Path, required=True)
    parser.add_argument("--segm-predictions", type=Path, required=True)
    parser.add_argument("--method", required=True)
    parser.add_argument("--seed", type=int, required=True)
    parser.add_argument("--split-role", choices=("development", "final_test"), required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--max-dets", type=int, nargs="+", default=(100, 300, 1000, 3000))
    parser.add_argument(
        "--expected-ap-at-max",
        type=float,
        required=True,
        help="Existing formal AP at the largest cap; exact-protocol equivalence gate.",
    )
    parser.add_argument("--ap-tolerance", type=float, default=1e-12)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    from pycocotools.coco import COCO

    for path in (args.ann_file, args.segm_predictions):
        if not path.is_file():
            raise FileNotFoundError(path)
    requested = sorted(set(args.max_dets))
    if len(requested) != len(args.max_dets) or any(value <= 10 for value in requested):
        raise ValueError("--max-dets values must be unique integers greater than 10")
    if args.ap_tolerance < 0:
        raise ValueError("--ap-tolerance must be non-negative")
    args.output_dir.mkdir(parents=True, exist_ok=True)

    annotation = load_json(args.ann_file)
    predictions = load_json(args.segm_predictions)
    image_ids = sorted(int(image["id"]) for image in annotation.get("images", []))
    if not image_ids:
        raise RuntimeError("Annotation contains no images")
    valid_ids = set(image_ids)
    unknown = sorted({int(record["image_id"]) for record in predictions} - valid_ids)
    if unknown:
        raise RuntimeError(f"Predictions contain unknown image IDs: {unknown[:8]}")

    counts = Counter(int(ann["image_id"]) for ann in annotation.get("annotations", []))
    coco_gt = COCO(str(args.ann_file))
    coco_dt = coco_gt.loadRes(predictions)
    metrics = run_coco_eval_multi_max(coco_gt, coco_dt, "segm", requested, image_ids)

    largest_cap = max(requested)
    reproduced_ap = float(metrics[largest_cap]["AP"])
    absolute_error = abs(reproduced_ap - args.expected_ap_at_max)
    if not math.isfinite(reproduced_ap) or absolute_error > args.ap_tolerance:
        raise RuntimeError(
            "Formal AP reproduction gate failed: "
            f"cap={largest_cap} expected={args.expected_ap_at_max:.17g} "
            f"observed={reproduced_ap:.17g} abs_error={absolute_error:.3g} "
            f"tolerance={args.ap_tolerance:.3g}"
        )

    rows = []
    for max_dets in requested:
        values = metrics[max_dets]
        row = {
            "method": args.method,
            "seed": args.seed,
            "split_role": args.split_role,
            "iou_type": "segm",
            "density": "all",
            "num_images": len(image_ids),
            "min_gt_instances": min(counts.get(image_id, 0) for image_id in image_ids),
            "max_gt_instances": max(counts.get(image_id, 0) for image_id in image_ids),
            "max_dets": max_dets,
            "AP": float(values["AP"]),
            "AP50": float(values["AP50"]),
            "AR": float(values["AR"]),
            "ARD": 1.0 - float(values["AR"]),
        }
        if not all(math.isfinite(row[key]) for key in ("AP", "AP50", "AR", "ARD")):
            raise RuntimeError(f"Non-finite metrics: {row}")
        rows.append(row)

    csv_path = args.output_dir / f"{args.method}_seed{args.seed}_maxdets_only.csv"
    fields = (
        "method", "seed", "split_role", "iou_type", "density", "num_images",
        "min_gt_instances", "max_gt_instances", "max_dets", "AP", "AP50", "AR", "ARD",
    )
    with csv_path.open("w", encoding="utf-8-sig", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)

    audit = {
        "status": "ok",
        "evaluation_scope": "segm_all_images_only",
        "method": args.method,
        "seed": args.seed,
        "split_role": args.split_role,
        "annotation_file": str(args.ann_file.resolve()),
        "annotation_sha256": sha256(args.ann_file),
        "segm_prediction_sha256": sha256(args.segm_predictions),
        "max_dets": requested,
        "num_images": len(image_ids),
        "rows": len(rows),
        "multi_max_dets_optimization": True,
        "formal_reproduction_gate": {
            "cap": largest_cap,
            "expected_AP": args.expected_ap_at_max,
            "observed_AP": reproduced_ap,
            "absolute_error": absolute_error,
            "tolerance": args.ap_tolerance,
            "passed": True,
        },
        "csv": str(csv_path.resolve()),
    }
    audit_path = args.output_dir / f"{args.method}_seed{args.seed}_maxdets_only_audit.json"
    audit_path.write_text(json.dumps(audit, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(audit, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()
