"""Evaluate a trained framework checkpoint on a chosen dataset/split, with the full
metric set INCLUDING boundary metrics. Two uses, one script (no retraining):

  * C7 (boundary fidelity, in-domain): point --eval_dataset at the same dataset to
    recompute boundary-Dice / NSD from a saved best.pth.
  * C4 (cross-center): set --eval_dataset/--eval_protocol to a DIFFERENT but
    label-compatible dataset (e.g. train busi? no — train cvc_clinicdb, eval
    kvasir_seg; both binary polyp). num_classes/in_channels must match.

Run from project root (…/NPJ), env seggen:
  CUDA_VISIBLE_DEVICES=5 python -m framework.eval_ckpt \
      --ckpt results/baselines/cvc_clinicdb_official/unet/seed0/best.pth \
      --arch unet --encoder resnet50 \
      --data_root $DR --dataset cvc_clinicdb --protocol official \
      --eval_dataset kvasir_seg --eval_protocol official \
      --out_json results/crosscenter/cvc2kvasir_unet_seed0.json
"""
from __future__ import annotations

import argparse
import json
import os
import sys

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

import numpy as np
import torch
from torch.utils.data import DataLoader

from framework.config import Config  # noqa: F401  (kept for parity / future YAML use)
from framework.models.registry import build_model, required_img_size
from framework.data.unified_dataset import UnifiedSegDataset
from framework.data.transforms import build_transform
from framework.metrics.metrics import per_image_metrics
from framework.metrics.boundary import boundary_metrics


def get_args():
    p = argparse.ArgumentParser("Evaluate a checkpoint (in-domain or cross-center) + boundary metrics")
    p.add_argument("--ckpt", required=True)
    p.add_argument("--arch", default="unet")
    p.add_argument("--encoder", default="resnet50")
    p.add_argument("--data_root", required=True)
    p.add_argument("--dataset", required=True, help="dataset the ckpt was TRAINED on (sets in_ch/num_classes)")
    p.add_argument("--protocol", required=True)
    p.add_argument("--eval_dataset", default="", help="dataset to evaluate ON (default = --dataset)")
    p.add_argument("--eval_protocol", default="", help="default = --protocol")
    p.add_argument("--split", default="test")
    p.add_argument("--img_size", type=int, default=256)
    p.add_argument("--normalize", default="auto")
    p.add_argument("--tol", type=float, default=2.0)
    p.add_argument("--batch_size", type=int, default=16)
    p.add_argument("--num_workers", type=int, default=6)
    p.add_argument("--out_json", default="")
    return p.parse_args()


def main():
    a = get_args()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    img_size = required_img_size(a.arch) or a.img_size
    eval_ds_name = a.eval_dataset or a.dataset
    eval_proto = a.eval_protocol or a.protocol

    # in_ch / num_classes are fixed by the TRAIN dataset (how the model was built)
    train_probe = UnifiedSegDataset(a.data_root, a.dataset, a.protocol, "test", transform=None)
    in_ch, n_cls = train_probe.in_channels, train_probe.num_classes

    model = build_model(a.arch, in_channels=in_ch, num_classes=n_cls, img_size=img_size,
                        encoder=a.encoder, encoder_weights="none", pretrained_ckpt="")
    ckpt = torch.load(a.ckpt, map_location="cpu", weights_only=False)
    model.load_state_dict(ckpt["model"])
    model.to(device).eval()

    tf = build_transform(img_size, in_ch, train=False, aug="none", normalize=a.normalize)
    ds = UnifiedSegDataset(a.data_root, eval_ds_name, eval_proto, a.split,
                           transform=tf, in_channels=in_ch, num_classes=n_cls)
    if ds.num_classes != n_cls:
        raise ValueError(f"label mismatch: train num_classes={n_cls} vs eval={ds.num_classes} "
                         f"({a.dataset}->{eval_ds_name}) — not label-compatible for cross-center.")
    cross = (eval_ds_name != a.dataset) or (eval_proto != a.protocol)
    print(f"[eval] ckpt={os.path.basename(a.ckpt)} train={a.dataset}/{a.protocol} "
          f"eval={eval_ds_name}/{eval_proto}/{a.split} cross_center={cross} "
          f"in_ch={in_ch} num_classes={n_cls} n={len(ds)}", flush=True)

    loader = DataLoader(ds, batch_size=a.batch_size, shuffle=False, num_workers=a.num_workers)
    recs = []
    for batch in loader:
        img = batch["image"].to(device, non_blocking=True)
        msk = batch["mask"].numpy()
        with torch.no_grad():
            pred = model(img).argmax(1).cpu().numpy()
        for i in range(pred.shape[0]):
            m = per_image_metrics(pred[i], msk[i], n_cls,
                                  include_background=False, compute_hd95=True)
            b = boundary_metrics(pred[i], msk[i], n_cls, tol=a.tol)
            m.update(b)
            recs.append(m)

    keys = [k for k, val in recs[0].items() if isinstance(val, (int, float))]  # skip per_class dict
    summary = {}
    for k in keys:
        v = np.array([r[k] for r in recs], dtype=np.float64)
        v = v[~np.isnan(v)]
        summary[k] = {"mean": round(float(v.mean()), 4) if v.size else None,
                      "std": round(float(v.std()), 4) if v.size else None, "n": int(v.size)}
    out = {"ckpt": a.ckpt, "train": f"{a.dataset}/{a.protocol}",
           "eval": f"{eval_ds_name}/{eval_proto}/{a.split}", "cross_center": cross,
           "num_images": len(ds), "metrics": summary}
    print(json.dumps(out, indent=2), flush=True)
    if a.out_json:
        os.makedirs(os.path.dirname(os.path.abspath(a.out_json)) or ".", exist_ok=True)
        with open(a.out_json, "w") as f:
            json.dump(out, f, indent=2)
        print(f"[eval] wrote {a.out_json}", flush=True)


if __name__ == "__main__":
    main()