#!/usr/bin/env python3
"""Recompute COCO metrics across maxDets and GT-density strata.

The script operates on saved prediction JSON files, so no GPU or retraining is
needed. Density groups are fixed a priori by ground-truth instances per image:
low <=100, medium 101-300, high 301-1000, extreme >1000.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
import multiprocessing
import sys
from collections import Counter
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path


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


DENSITY_RULES = (
    ("low", 0, 100),
    ("medium", 101, 300),
    ("high", 301, 1000),
    ("extreme", 1001, math.inf),
)


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("--bbox-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(
        "--workers",
        type=int,
        default=1,
        help="Parallel COCOeval workers. Linux fork shares loaded predictions copy-on-write.",
    )
    return parser.parse_args()


def load_json(path: Path):
    with path.open(encoding="utf-8") as handle:
        return json.load(handle)


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def density_groups(annotation: dict) -> tuple[dict[str, list[int]], dict[int, int]]:
    counts = Counter(int(ann["image_id"]) for ann in annotation.get("annotations", []))
    image_ids = [int(image["id"]) for image in annotation.get("images", [])]
    groups = {name: [] for name, _, _ in DENSITY_RULES}
    for image_id in image_ids:
        count = counts.get(image_id, 0)
        for name, lower, upper in DENSITY_RULES:
            if lower <= count <= upper:
                groups[name].append(image_id)
                break
    return groups, {image_id: counts.get(image_id, 0) for image_id in image_ids}


def _mean_valid(values) -> float:
    values = values[values > -1]
    return float(values.mean()) if values.size else -1.0


def _metrics_at_max_dets(evaluator, max_dets: int) -> dict[str, float]:
    """Extract AP/AP50/AR for one maxDets value from one COCOeval pass."""
    import numpy as np

    params = evaluator.params
    area_index = params.areaRngLbl.index("all")
    max_dets_index = params.maxDets.index(max_dets)
    precision = evaluator.eval["precision"][:, :, :, area_index, max_dets_index]
    recall = evaluator.eval["recall"][:, :, area_index, max_dets_index]
    iou50 = np.where(np.isclose(params.iouThrs, 0.50))[0]
    if not iou50.size:
        raise RuntimeError("COCOeval parameters do not contain IoU=0.50")
    return {
        "AP": _mean_valid(precision),
        "AP50": _mean_valid(precision[iou50]),
        "AR": _mean_valid(recall),
    }


def run_coco_eval_multi_max(
    coco_gt,
    coco_dt,
    iou_type: str,
    max_dets_values: list[int],
    image_ids: list[int],
) -> dict[int, dict[str, float]]:
    """Evaluate once at the largest maxDets and extract every requested cap.

    COCOeval computes image-level matches at ``params.maxDets[-1]``. During
    accumulation it slices those ranked detections for every entry in
    ``params.maxDets``. Supplying all caps together is therefore equivalent to
    separate runs, while avoiding repeated O(GT x detections) matching.
    """
    from pycocotools.cocoeval import COCOeval
    from livecell_protocol import LIVECELL_AREA_RANGE_LABELS, LIVECELL_AREA_RANGES

    requested = sorted(set(max_dets_values))
    evaluator = COCOeval(coco_gt, coco_dt, iouType=iou_type)
    evaluator.params.maxDets = sorted({1, 10, *requested})
    evaluator.params.areaRng = LIVECELL_AREA_RANGES
    evaluator.params.areaRngLbl = LIVECELL_AREA_RANGE_LABELS
    evaluator.params.imgIds = sorted(set(image_ids))
    evaluator.evaluate()
    evaluator.accumulate()
    return {max_dets: _metrics_at_max_dets(evaluator, max_dets) for max_dets in requested}


_WORKER_COCO_GT = None
_WORKER_COCO_DT = None


def _evaluate_stratum(job):
    iou_type, density, image_ids, max_dets_values = job
    if _WORKER_COCO_GT is None or _WORKER_COCO_DT is None:
        raise RuntimeError("COCOeval worker state was not initialized")
    metrics = run_coco_eval_multi_max(
        _WORKER_COCO_GT,
        _WORKER_COCO_DT[iou_type],
        iou_type,
        max_dets_values,
        image_ids,
    )
    return iou_type, density, metrics


def main() -> None:
    args = parse_args()
    # Keep this import inside main so density stratification can be unit-tested
    # without loading torch/torchvision on analysis-only machines.
    from pycocotools.coco import COCO
    for path in (args.ann_file, args.segm_predictions, args.bbox_predictions):
        if not path.is_file():
            raise FileNotFoundError(path)
    if len(set(args.max_dets)) != len(args.max_dets) or any(value <= 10 for value in args.max_dets):
        raise ValueError("--max-dets values must be unique integers greater than 10")
    if args.workers <= 0:
        raise ValueError("--workers must be positive")
    args.output_dir.mkdir(parents=True, exist_ok=True)

    annotation = load_json(args.ann_file)
    groups, counts = density_groups(annotation)
    predictions = {
        "bbox": load_json(args.bbox_predictions),
        "segm": load_json(args.segm_predictions),
    }
    valid_image_ids = set(counts)
    for iou_type, records in predictions.items():
        unknown = sorted({int(record["image_id"]) for record in records} - valid_image_ids)
        if unknown:
            raise RuntimeError(f"{iou_type} predictions contain unknown image IDs: {unknown[:8]}")

    rows = []
    strata = [("all", sorted(valid_image_ids))] + [(name, ids) for name, ids in groups.items()]
    nonempty_strata = [(density, image_ids) for density, image_ids in strata if image_ids]
    coco_gt = COCO(str(args.ann_file))
    coco_dt = {iou_type: coco_gt.loadRes(records) for iou_type, records in predictions.items()}
    jobs = [
        (iou_type, density, image_ids, list(args.max_dets))
        for iou_type in predictions
        for density, image_ids in nonempty_strata
    ]

    global _WORKER_COCO_GT, _WORKER_COCO_DT
    _WORKER_COCO_GT = coco_gt
    _WORKER_COCO_DT = coco_dt
    worker_count = min(args.workers, len(jobs))
    if worker_count == 1:
        evaluated = [_evaluate_stratum(job) for job in jobs]
    else:
        context = multiprocessing.get_context("fork")
        with ProcessPoolExecutor(max_workers=worker_count, mp_context=context) as executor:
            evaluated = list(executor.map(_evaluate_stratum, jobs))
    metrics_by_stratum = {
        (iou_type, density): metrics
        for iou_type, density, metrics in evaluated
    }

    # Preserve the legacy CSV order so old and optimized outputs are directly diffable.
    for iou_type in predictions:
        for max_dets in sorted(args.max_dets):
            for density, image_ids in nonempty_strata:
                metrics = metrics_by_stratum[(iou_type, density)][max_dets]
                ar = float(metrics["AR"])
                row = {
                    "method": args.method,
                    "seed": args.seed,
                    "split_role": args.split_role,
                    "iou_type": iou_type,
                    "density": density,
                    "num_images": len(image_ids),
                    "min_gt_instances": min(counts[image_id] for image_id in image_ids),
                    "max_gt_instances": max(counts[image_id] for image_id in image_ids),
                    "max_dets": max_dets,
                    "AP": float(metrics["AP"]),
                    "AP50": float(metrics["AP50"]),
                    "AR": ar,
                    "ARD": 1.0 - 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}_dense_protocol_ablation.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",
        "method": args.method,
        "seed": args.seed,
        "split_role": args.split_role,
        "annotation_file": str(args.ann_file.resolve()),
        "annotation_sha256": sha256(args.ann_file),
        "prediction_sha256": {
            "bbox": sha256(args.bbox_predictions),
            "segm": sha256(args.segm_predictions),
        },
        "density_rules": [
            {"name": name, "min_inclusive": lower, "max_inclusive": None if math.isinf(upper) else upper}
            for name, lower, upper in DENSITY_RULES
        ],
        "density_image_counts": {name: len(ids) for name, ids in groups.items()},
        "max_dets": sorted(args.max_dets),
        "rows": len(rows),
        "evaluation_passes": len(jobs),
        "workers": worker_count,
        "multi_max_dets_optimization": True,
        "csv": str(csv_path.resolve()),
    }
    json_path = args.output_dir / f"{args.method}_seed{args.seed}_dense_protocol_ablation_audit.json"
    json_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()
