"""
evaluate.py
============
توابع ماژولار ارزیابی روی مجموعه تست (Image-AUROC, Pixel-AUROC, Pixel-AP).
"""

import os
import sys
import numpy as np
import torch
import torch.nn.functional as F
from tqdm.auto import tqdm

sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
from data.data_loader import list_test_images, load_and_resize, load_mask
from utils.metrics import anomaly_score_from_seg, compute_image_auroc, compute_pixel_auroc, compute_pixel_ap


@torch.no_grad()
def evaluate_category(model, root: str, category: str, size, device: str, dataset_name: str = "mvtec",
                      smooth_kernel: int = 7, smooth_sigma: float = 1.5):
    """
    Run model on one class test set. Returns image/pixel scores and labels.
    """
    model.eval()
    test_items = list_test_images(root, category, dataset_name)

    image_scores, image_labels = [], []
    pixel_scores, pixel_labels = [], []

    for img_path, defect_type, mask_path in test_items:
        img = load_and_resize(img_path, size)
        img_t = torch.from_numpy(img.transpose(2, 0, 1)).float().unsqueeze(0).to(device)

        # Eval in FP32 — AMP leftovers in weights can still NaN; sanitize scores below
        with torch.amp.autocast("cuda", enabled=False):
            out = model(img_t)
        seg_map = anomaly_score_from_seg(
            out["seg_logits"].float(), smooth_kernel=smooth_kernel, smooth_sigma=smooth_sigma
        )[0].cpu().numpy()  # (H,W)
        seg_map = np.nan_to_num(seg_map, nan=0.0, posinf=1.0, neginf=0.0)

        label = 0 if defect_type == "good" else 1
        image_scores.append(float(np.max(seg_map)))
        image_labels.append(label)

        if mask_path is not None:
            gt = load_mask(mask_path, size)
        else:
            gt = np.zeros(size, dtype=np.float32)

        pixel_scores.append(seg_map)
        pixel_labels.append(gt)

    return {
        "image_scores": np.asarray(image_scores, dtype=np.float64),
        "image_labels": np.asarray(image_labels),
        "pixel_scores": np.stack(pixel_scores, axis=0),
        "pixel_labels": np.stack(pixel_labels, axis=0),
    }


def evaluate_dataset(model, cfg: dict, dataset_name: str = "mvtec", device: str = "cuda",
                     verbose: bool = True):
    """
    Full eval over all classes. Returns dict with per-category metrics + "average".
    Set verbose=False to silence per-category prints (training loop uses this).
    """
    if dataset_name == "mvtec":
        root = cfg["data"]["mvtec_root"]
        categories = cfg["data"]["mvtec_categories"]
        size = tuple(cfg["data"]["mvtec_size"])
    else:
        root = cfg["data"]["visa_root"]
        categories = cfg["data"]["visa_categories"]
        size = tuple(cfg["data"]["visa_size"])

    smooth_k = cfg["model"].get("anomaly_smooth_kernel", 7)
    smooth_s = cfg["model"].get("anomaly_smooth_sigma", 1.5)

    results = {}
    cat_iter = categories if verbose else tqdm(categories, desc="eval", leave=False, dynamic_ncols=True)
    for cat in cat_iter:
        try:
            res = evaluate_category(
                model, root, cat, size, device, dataset_name,
                smooth_kernel=smooth_k, smooth_sigma=smooth_s,
            )
        except Exception as e:
            print(f"[WARN] {cat}: eval failed ({e}); skipping.")
            continue
        if len(np.unique(res["image_labels"])) < 2:
            if verbose:
                print(f"[WARN] {cat}: need both normal/anomaly labels for AUROC.")
            continue
        i_auroc = compute_image_auroc(res["image_scores"], res["image_labels"])
        p_auroc = compute_pixel_auroc(res["pixel_scores"], res["pixel_labels"])
        p_ap = compute_pixel_ap(res["pixel_scores"], res["pixel_labels"])
        if not np.isfinite(i_auroc) or not np.isfinite(p_auroc):
            print(f"[WARN] {cat}: non-finite metrics (I={i_auroc}, P={p_auroc}); skipping.")
            continue
        results[cat] = {"I-AUROC": i_auroc, "P-AUROC": p_auroc, "P-AP": p_ap}
        if verbose:
            print(f"{cat:15s} | I-AUROC={i_auroc:.3f}  P-AUROC={p_auroc:.3f}  P-AP={p_ap:.3f}")

    if results:
        avg = {k: float(np.mean([v[k] for v in results.values()])) for k in ["I-AUROC", "P-AUROC", "P-AP"]}
        results["average"] = avg
        if verbose:
            print(f"{'AVERAGE':15s} | I-AUROC={avg['I-AUROC']:.3f}  P-AUROC={avg['P-AUROC']:.3f}  P-AP={avg['P-AP']:.3f}")
    else:
        results["average"] = {"I-AUROC": 0.0, "P-AUROC": 0.0, "P-AP": 0.0}
        print("[WARN] No valid categories in eval — returning zero averages.")

    return results