"""Compute framework-format metrics.json from nnU-Net / U-Mamba test predictions.

nnU-Net only reports validation Dice during training. To compare on the SAME
held-out test set with the SAME 7 metrics as the framework, we: predict on
imagesTs (done separately via nnUNetv2_predict), then run THIS script to score the
predicted masks against labelsTs using framework/metrics.py, writing a metrics.json
in the exact framework format so report/aggregate.py includes nnU-Net/U-Mamba rows.

  python framework/nnunet_eval.py --data_root <processed_unified> --dataset <ds> \
     --protocol <proto> --raw <nnUNet_raw> --dataset_id <ID> --fold <f> \
     --pred_dir <predictions> --arch nnunet --exp_name baselines
"""
from __future__ import annotations

import os
import sys
import json
import glob
import argparse

import numpy as np
import cv2

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from framework.metrics.metrics import per_image_metrics, aggregate
from framework.data.unified_dataset import _read_metadata, detect_num_classes


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data_root", required=True)
    ap.add_argument("--dataset", required=True)
    ap.add_argument("--protocol", required=True)
    ap.add_argument("--raw", required=True, help="nnUNet_raw root")
    ap.add_argument("--dataset_id", type=int, required=True)
    ap.add_argument("--name", default="", help="Dataset<DDD>_<name> suffix; default <dataset>_<protocol>")
    ap.add_argument("--fold", type=int, required=True)
    ap.add_argument("--pred_dir", required=True)
    ap.add_argument("--arch", default="nnunet")
    ap.add_argument("--exp_name", default="baselines")
    ap.add_argument("--out_root", default="results")
    ap.add_argument("--include_background", action="store_true")
    ap.add_argument("--eval_size", type=int, default=0,
                    help="resize pred+gt to R×R (nearest) before scoring; 0 = native GT resolution")
    args = ap.parse_args()

    name = args.name or f"{args.dataset}_{args.protocol}"
    dsname = f"Dataset{args.dataset_id:03d}_{name}"
    lab_dir = os.path.join(args.raw, dsname, "labelsTs")

    meta = _read_metadata(args.data_root, args.dataset)
    gt_masks = sorted(glob.glob(os.path.join(lab_dir, "*.png")))
    num_classes = detect_num_classes(meta, gt_masks, args.dataset)

    records = []
    n_missing = 0
    for pp in sorted(glob.glob(os.path.join(args.pred_dir, "*.png"))):
        base = os.path.basename(pp)
        gp = os.path.join(lab_dir, base)
        if not os.path.isfile(gp):
            n_missing += 1
            continue
        pred = cv2.imread(pp, cv2.IMREAD_GRAYSCALE)
        gt = cv2.imread(gp, cv2.IMREAD_GRAYSCALE)
        if pred is None or gt is None:
            n_missing += 1
            continue
        if pred.shape != gt.shape:
            pred = cv2.resize(pred, (gt.shape[1], gt.shape[0]), interpolation=cv2.INTER_NEAREST)
        if args.eval_size > 0:  # resolution-fair common-R scoring
            R = args.eval_size
            pred = cv2.resize(pred, (R, R), interpolation=cv2.INTER_NEAREST)
            gt = cv2.resize(gt, (R, R), interpolation=cv2.INTER_NEAREST)
        records.append(per_image_metrics(pred.astype(np.int64), gt.astype(np.int64),
                                         num_classes, include_background=args.include_background,
                                         compute_hd95=True))
    if not records:
        raise SystemExit(f"no matched (pred,gt) pairs in {args.pred_dir} vs {lab_dir}")

    agg = aggregate(records)
    out = {"dataset": args.dataset, "protocol": args.protocol, "arch": args.arch,
           "seed": args.fold, "num_classes": num_classes, "metrics": agg, "per_image": records}
    out_dir = os.path.join(args.out_root, args.exp_name, f"{args.dataset}_{args.protocol}",
                           args.arch, f"seed{args.fold}")
    os.makedirs(out_dir, exist_ok=True)
    with open(os.path.join(out_dir, "metrics.json"), "w") as f:
        json.dump(out, f, indent=2)
    print(f"[nnunet_eval] {dsname} fold{args.fold}: n={len(records)} (missing {n_missing}) "
          f"dice={agg['dice_mean']:.4f} iou={agg['iou_mean']:.4f} hd95={agg['hd95_mean']:.2f} "
          f"-> {out_dir}/metrics.json")


if __name__ == "__main__":
    main()