#!/usr/bin/env python3
"""Evaluate a completed dense-protocol checkpoint on an explicit COCO split.

This is inference only. It never updates the checkpoint and requires an explicit
annotation file so that development validation and final test evaluation cannot
be confused accidentally.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import sys
from pathlib import Path

import torch


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

from train_maskrcnn_baseline import (  # noqa: E402
    LiveCellCocoDataset,
    evaluate_and_visualize,
)
from train_maskrcnn_smoke import build_model, index_images  # noqa: E402


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-root", type=Path, required=True)
    parser.add_argument("--ann-file", type=Path, required=True)
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--split-role", choices=("development", "final_test"), required=True)
    parser.add_argument("--device", choices=("auto", "cuda", "cpu"), default="auto")
    parser.add_argument("--max-images", type=int, default=0)
    parser.add_argument("--vis-count", type=int, default=0)
    parser.add_argument(
        "--defer-coco-eval",
        action="store_true",
        help="Save inference artifacts and defer metrics to reevaluate_dense_predictions.py.",
    )
    return parser.parse_args()


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 number(config: dict, key: str, default, cast):
    value = config.get("args", {}).get(key, default)
    return cast(value)


def main() -> None:
    args = parse_args()
    for path in (args.data_root, args.ann_file, args.checkpoint):
        if not path.exists():
            raise FileNotFoundError(path)
    args.output_dir.mkdir(parents=True, exist_ok=True)
    if any(args.output_dir.iterdir()):
        raise RuntimeError(f"Output directory must be empty: {args.output_dir}")

    device = torch.device(
        "cuda" if args.device == "auto" and torch.cuda.is_available() else
        "cpu" if args.device == "auto" else args.device
    )
    checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    state = checkpoint["model"] if isinstance(checkpoint, dict) and "model" in checkpoint else checkpoint
    config = checkpoint.get("config", {}) if isinstance(checkpoint, dict) else {}

    model = build_model(
        num_classes=2,
        min_size=number(config, "min_size", 512, int),
        max_size=number(config, "max_size", 768, int),
        box_detections_per_img=number(config, "box_detections_per_image", 3000, int),
        rpn_pre_nms_top_n_train=number(config, "rpn_pre_nms_top_n_train", 6000, int),
        rpn_post_nms_top_n_train=number(config, "rpn_post_nms_top_n_train", 3000, int),
        rpn_pre_nms_top_n_test=number(config, "rpn_pre_nms_top_n_test", 6000, int),
        rpn_post_nms_top_n_test=number(config, "rpn_post_nms_top_n_test", 3000, int),
        box_batch_size_per_image=number(config, "box_batch_size_per_image", 512, int),
    )
    missing, unexpected = model.load_state_dict(state, strict=False)
    if missing or unexpected:
        raise RuntimeError(
            f"Checkpoint/model mismatch: missing={missing[:8]}, unexpected={unexpected[:8]}"
        )
    model.to(device)

    dataset = LiveCellCocoDataset(
        data_root=args.data_root,
        annotation_file=args.ann_file,
        image_index=index_images(args.data_root),
        max_images=args.max_images,
        max_instances_per_image=0,
        seed=104729,
        training=False,
        crop_size=0,
    )
    result = evaluate_and_visualize(
        model=model,
        dataset=dataset,
        device=device,
        output_dir=args.output_dir,
        score_threshold=number(config, "eval_score_threshold", 0.05, float),
        max_predictions=number(config, "eval_max_predictions", 3000, int),
        eval_max_images=args.max_images,
        vis_count=args.vis_count,
        use_livecell_area_ranges=True,
        defer_coco_eval=args.defer_coco_eval,
    )
    if args.defer_coco_eval:
        prediction_files = {
            "bbox": args.output_dir / "predictions_bbox.json",
            "segm": args.output_dir / "predictions_segm.json",
            "quick_eval": args.output_dir / "quick_eval.csv",
        }
        for path in prediction_files.values():
            if not path.is_file() or path.stat().st_size == 0:
                raise RuntimeError(f"Missing or empty inference artifact: {path}")
        audit = {
            "status": "ok",
            "metrics_status": "deferred_to_dense_protocol_evaluator",
            "split_role": args.split_role,
            "annotation_file": str(args.ann_file.resolve()),
            "annotation_sha256": sha256(args.ann_file),
            "checkpoint": str(args.checkpoint.resolve()),
            "checkpoint_sha256": sha256(args.checkpoint),
            "device": str(device),
            "dataset": dataset.describe(),
            "inference_summary": result,
            "artifact_sha256": {
                name: sha256(path) for name, path in prediction_files.items()
            },
        }
        output = args.output_dir / "inference_complete.json"
        output.write_text(json.dumps(audit, ensure_ascii=False, indent=2), encoding="utf-8")
        print(json.dumps(audit, ensure_ascii=False, indent=2))
        return

    for iou_type in ("bbox", "segm"):
        metrics = result[iou_type]
        ar = float(metrics["AR_max"])
        metrics["ARD"] = 1.0 - ar
        metrics.pop("AFNR", None)
        if not all(math.isfinite(float(metrics[key])) for key in ("AP", "AP50", "AR_max", "ARD")):
            raise RuntimeError(f"Non-finite {iou_type} metrics: {metrics}")

    audit = {
        "status": "ok",
        "split_role": args.split_role,
        "annotation_file": str(args.ann_file.resolve()),
        "annotation_sha256": sha256(args.ann_file),
        "checkpoint": str(args.checkpoint.resolve()),
        "checkpoint_sha256": sha256(args.checkpoint),
        "device": str(device),
        "dataset": dataset.describe(),
        "eval_summary": result,
    }
    output = args.output_dir / "final_evaluation_summary.json"
    output.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()
