"""
生成 mask 群体级评估
=====================
四个维度:

1. Mask FID + Density (Fidelity)
   - 训一个小卷积 autoencoder 编码 mask -> 64维 latent
   - FID: 两组 latent 的 Fréchet 距离 (越小越真)
   - Density (Naeem 2020): 生成样本周围真实样本的局部密度 (越高越真)

2. Coverage (Diversity)
   - 真实样本中被生成样本覆盖的比例 (越高越多样)

3. 1D Shape Descriptors KDE 对比
   - area / aspect_ratio / circularity / solidity / extent / centroid_x / centroid_y
   - 每个描述子画 gen vs GT KDE; 算 Wasserstein 距离

4. Severity-conditional spread (群体层面)
   - 跨 severity 类间距离 / 同 severity 类内距离 (>1 说明 prompt 真的控制了形态)
   - 真实 vs 生成 的 severity 间差距方向是否一致

输入 (与之前 eval_gen_mask 一致):
  --jsonl, --gen_root, --gt_root, --out_dir

依赖: torch, scipy, scikit-learn, matplotlib
"""
import argparse
import json
import re
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from PIL import Image
from scipy.stats import wasserstein_distance
from sklearn.neighbors import NearestNeighbors


# ======================================================================
#  数据
# ======================================================================
def load_mask(path):
    arr = np.array(Image.open(path).convert("L"))
    return (arr > 127).astype(np.uint8)


def parse_prompt_severity(prompt):
    for s in ["mild", "moderate", "severe"]:
        if re.search(rf"\b{s}\b", prompt.lower()):
            return s
    return None


def collect_masks(jsonl, gen_root, gt_root):
    """返回 list of dict: {gen_mask, gt_mask, disease, severity, id}.
       生成/GT 任一缺失则跳过."""
    gen_root, gt_root = Path(gen_root), Path(gt_root)
    items = [json.loads(l) for l in open(jsonl) if l.strip()]
    out = []
    n_skip = 0
    for it in items:
        gen_p = gen_root / it["attn_list"][0][1]
        gt_p = gt_root / it["mask"]
        if not gen_p.is_file() or not gt_p.is_file():
            n_skip += 1
            continue
        out.append({
            "id": Path(it["file_name"]).parent.name,
            "disease": it["attn_list"][0][0],
            "severity": parse_prompt_severity(it["prompt"]),
            "gen": load_mask(gen_p),
            "gt": load_mask(gt_p),
        })
    print(f"Loaded {len(out)} pairs (skipped {n_skip})")
    return out


# ======================================================================
#  1. Mask AutoEncoder (训练 -> 提 latent)
# ======================================================================
class MaskAE(nn.Module):
    def __init__(self, latent_dim=64, in_size=128):
        super().__init__()
        c = 32
        # encoder
        self.enc = nn.Sequential(
            nn.Conv2d(1, c, 4, 2, 1), nn.ReLU(True),      # 64
            nn.Conv2d(c, c*2, 4, 2, 1), nn.ReLU(True),    # 32
            nn.Conv2d(c*2, c*4, 4, 2, 1), nn.ReLU(True),  # 16
            nn.Conv2d(c*4, c*8, 4, 2, 1), nn.ReLU(True),  # 8
            nn.Flatten(),
            nn.Linear(c*8 * 8 * 8, latent_dim),
        )
        # decoder
        self.dec_fc = nn.Linear(latent_dim, c*8 * 8 * 8)
        self.dec = nn.Sequential(
            nn.Unflatten(1, (c*8, 8, 8)),
            nn.ConvTranspose2d(c*8, c*4, 4, 2, 1), nn.ReLU(True),
            nn.ConvTranspose2d(c*4, c*2, 4, 2, 1), nn.ReLU(True),
            nn.ConvTranspose2d(c*2, c, 4, 2, 1), nn.ReLU(True),
            nn.ConvTranspose2d(c, 1, 4, 2, 1),
        )

    def encode(self, x):
        return self.enc(x)

    def forward(self, x):
        z = self.encode(x)
        return self.dec(self.dec_fc(z))


def train_ae(real_masks, latent_dim=64, ae_size=128, epochs=20, bs=64, device="cuda"):
    """real_masks: list of np.uint8 (H,W). 返回 trained AE."""
    print(f"\n[AE] training on {len(real_masks)} real masks "
          f"({epochs} epochs, latent={latent_dim})")
    # 预处理: resize + to tensor
    tensors = []
    for m in real_masks:
        im = Image.fromarray(m * 255).resize((ae_size, ae_size), Image.NEAREST)
        tensors.append(np.array(im, dtype=np.float32) / 255.0)
    x = torch.from_numpy(np.stack(tensors))[:, None]  # [N,1,H,W]

    ae = MaskAE(latent_dim, ae_size).to(device)
    opt = torch.optim.AdamW(ae.parameters(), lr=1e-3)
    loss_fn = nn.BCEWithLogitsLoss()

    ae.train()
    for ep in range(epochs):
        # 随机打乱
        idx = torch.randperm(len(x))
        total = 0
        for i in range(0, len(x), bs):
            batch = x[idx[i:i+bs]].to(device)
            logits = ae(batch)
            loss = loss_fn(logits, batch)
            opt.zero_grad(); loss.backward(); opt.step()
            total += loss.item() * len(batch)
        print(f"  ep{ep+1:02d}  bce={total/len(x):.4f}")
    return ae


@torch.no_grad()
def encode_masks(ae, masks, ae_size=128, bs=128, device="cuda"):
    ae.eval()
    feats = []
    for i in range(0, len(masks), bs):
        chunk = []
        for m in masks[i:i+bs]:
            im = Image.fromarray(m * 255).resize((ae_size, ae_size), Image.NEAREST)
            chunk.append(np.array(im, dtype=np.float32) / 255.0)
        x = torch.from_numpy(np.stack(chunk))[:, None].to(device)
        z = ae.encode(x).cpu().numpy()
        feats.append(z)
    return np.concatenate(feats, 0)


# ======================================================================
#  2. FID / Density / Coverage
# ======================================================================
def frechet_distance(mu1, sig1, mu2, sig2):
    from scipy.linalg import sqrtm
    diff = mu1 - mu2
    covmean = sqrtm(sig1 @ sig2)
    if np.iscomplexobj(covmean):
        covmean = covmean.real
    return float(diff @ diff + np.trace(sig1 + sig2 - 2 * covmean))


def compute_fid(feat_gen, feat_real):
    mu_g, sig_g = feat_gen.mean(0), np.cov(feat_gen, rowvar=False)
    mu_r, sig_r = feat_real.mean(0), np.cov(feat_real, rowvar=False)
    return frechet_distance(mu_g, sig_g, mu_r, sig_r)


def density_coverage(feat_gen, feat_real, k=5):
    """Naeem et al. 2020. 返回 (density, coverage)."""
    # 真实样本的 k-NN 半径
    nn_real = NearestNeighbors(n_neighbors=k+1).fit(feat_real)
    real_radii = nn_real.kneighbors(feat_real)[0][:, -1]   # 每点到第 k 邻的距离
    # 生成样本的 k-NN 半径
    nn_gen = NearestNeighbors(n_neighbors=k+1).fit(feat_gen)
    gen_radii = nn_gen.kneighbors(feat_gen)[0][:, -1]

    # density: 每个生成点落在多少个真实球内 / k
    dist_g2r = NearestNeighbors().fit(feat_real).kneighbors(
        feat_gen, n_neighbors=len(feat_real), return_distance=True)[0]
    in_balls = (dist_g2r < real_radii[None, :]).sum(1)
    density = float(in_balls.mean() / k)

    # coverage: 多少真实点在它最近生成点的 gen_radius 之内
    nn_r2g = NearestNeighbors(n_neighbors=1).fit(feat_gen)
    d_r2g, idx_r2g = nn_r2g.kneighbors(feat_real)
    covered = (d_r2g[:, 0] < gen_radii[idx_r2g[:, 0]]).mean()
    return density, float(covered)


# ======================================================================
#  3. Shape descriptors
# ======================================================================
def shape_descriptors(mask):
    """返回 dict of 7 描述子. 全空 mask 返回 NaN."""
    if mask.sum() == 0:
        return {k: np.nan for k in
                ["area", "aspect_ratio", "circularity", "solidity",
                 "extent", "cx", "cy"]}
    H, W = mask.shape
    ys, xs = np.where(mask > 0)
    area = int(len(ys))
    h = ys.max() - ys.min() + 1
    w = xs.max() - xs.min() + 1
    bbox_area = h * w
    extent = area / bbox_area              # 病灶 vs 包围盒
    aspect = h / w if w > 0 else np.nan

    # circularity = 4πA / P²   用 OpenCV-free 近似: 边界点数当周长
    # 用 numpy 简易法: gradient 非零点数当周长
    from scipy.ndimage import binary_erosion
    eroded = binary_erosion(mask)
    perimeter = int((mask & ~eroded).sum())
    circ = 4 * np.pi * area / (perimeter ** 2) if perimeter > 0 else np.nan

    # solidity = area / convex_hull_area
    try:
        from scipy.spatial import ConvexHull
        pts = np.column_stack([xs, ys])
        if len(pts) >= 3:
            hull = ConvexHull(pts)
            solidity = area / hull.volume   # 2D 时 volume = area
        else:
            solidity = 1.0
    except Exception:
        solidity = np.nan

    cx = (xs.mean() / W)   # 归一化重心
    cy = (ys.mean() / H)

    return {"area": area, "aspect_ratio": aspect, "circularity": circ,
            "solidity": solidity, "extent": extent, "cx": cx, "cy": cy}


def shape_descriptors_table(records, side):
    """side = 'gen' or 'gt'."""
    rows = []
    for r in records:
        d = shape_descriptors(r[side])
        d["disease"] = r["disease"]
        d["severity"] = r["severity"]
        d["source"] = side
        rows.append(d)
    return pd.DataFrame(rows)


def plot_descriptor_kde(df_gen, df_gt, out_dir, descriptors=None, by="disease"):
    """每个描述子一张图, 内含所有疾病 (gen vs gt 配对色)."""
    if descriptors is None:
        descriptors = ["area", "aspect_ratio", "circularity",
                       "solidity", "extent", "cx", "cy"]
    out_dir.mkdir(exist_ok=True, parents=True)
    diseases = sorted(set(df_gen["disease"]) | set(df_gt["disease"]))
    for d in descriptors:
        ncol = 4
        nrow = (len(diseases) + ncol - 1) // ncol
        fig, axes = plt.subplots(nrow, ncol, figsize=(ncol*4, nrow*3))
        axes = np.array(axes).flatten()
        for i, dis in enumerate(diseases):
            ax = axes[i]
            gen_v = df_gen[(df_gen["disease"]==dis)][d].dropna().values
            gt_v  = df_gt [(df_gt ["disease"]==dis)][d].dropna().values
            if len(gen_v) > 5 and len(gt_v) > 5:
                # log 化 area 让分布可读
                if d == "area":
                    gen_v = np.log10(gen_v + 1)
                    gt_v = np.log10(gt_v + 1)
                ax.hist(gt_v,  bins=30, alpha=0.5, label="GT",  density=True, color="#4DBBA1")
                ax.hist(gen_v, bins=30, alpha=0.5, label="Gen", density=True, color="#E89A6B")
                ax.set_title(f"{dis} (n={len(gen_v)})", fontsize=10)
                ax.tick_params(labelsize=8)
                if i == 0:
                    ax.legend(fontsize=8)
        for j in range(len(diseases), len(axes)):
            axes[j].axis("off")
        fig.suptitle(f"{d}{'  (log10)' if d=='area' else ''} — Gen vs GT", fontsize=14)
        fig.tight_layout()
        fig.savefig(out_dir / f"kde_{d}.png", dpi=130, bbox_inches="tight")
        plt.close(fig)


def shape_wasserstein(df_gen, df_gt, descriptors):
    """每 (disease, descriptor) 算 1D Wasserstein."""
    rows = []
    for dis in sorted(set(df_gen["disease"]) | set(df_gt["disease"])):
        g = df_gen[df_gen["disease"] == dis]
        t = df_gt [df_gt ["disease"] == dis]
        row = {"disease": dis, "n_gen": len(g), "n_gt": len(t)}
        for d in descriptors:
            gv = g[d].dropna().values
            tv = t[d].dropna().values
            if len(gv) >= 5 and len(tv) >= 5:
                # area 用 log
                if d == "area":
                    gv = np.log10(gv + 1); tv = np.log10(tv + 1)
                row[f"W_{d}"] = wasserstein_distance(gv, tv)
            else:
                row[f"W_{d}"] = np.nan
        rows.append(row)
    return pd.DataFrame(rows)


# ======================================================================
#  4. Severity-conditional spread (群体)
# ======================================================================
def latent_pairwise_dist(feats):
    """[N, D] -> 平均对 L2 距离."""
    if len(feats) < 2:
        return np.nan
    n = min(len(feats), 200)   # 大时下采样, 防 O(N²) 太慢
    if len(feats) > n:
        idx = np.random.choice(len(feats), n, replace=False)
        feats = feats[idx]
    from scipy.spatial.distance import pdist
    return float(pdist(feats).mean())


def severity_spread(records, ae, ae_size, device, side="gen"):
    """对每个 disease, 算同/跨 severity 的 latent 距离."""
    # 按 (disease, severity) 分桶
    buckets = {}
    for r in records:
        if r["severity"] is None:
            continue
        key = (r["disease"], r["severity"])
        buckets.setdefault(key, []).append(r[side])
    # 每桶编码
    bucket_feats = {}
    for key, masks in buckets.items():
        if len(masks) < 2:
            continue
        bucket_feats[key] = encode_masks(ae, masks, ae_size, device=device)

    diseases = sorted(set(d for d, _ in bucket_feats))
    rows = []
    for dis in diseases:
        within = []   # 同 severity 类内距离
        between = []  # 跨 severity 类间距离 (类中心间)
        sev_means = {}
        for sev in ["mild", "moderate", "severe"]:
            key = (dis, sev)
            if key in bucket_feats:
                within.append(latent_pairwise_dist(bucket_feats[key]))
                sev_means[sev] = bucket_feats[key].mean(0)
        if len(sev_means) < 2:
            continue
        # 类中心间 L2
        sev_keys = list(sev_means.keys())
        bc = []
        for i in range(len(sev_keys)):
            for j in range(i+1, len(sev_keys)):
                bc.append(np.linalg.norm(sev_means[sev_keys[i]] - sev_means[sev_keys[j]]))
        within_avg = np.nanmean(within)
        between_avg = np.mean(bc)
        row = {"disease": dis, "side": side,
               "within_severity": within_avg,
               "between_severity": between_avg,
               "between/within": between_avg / within_avg if within_avg > 0 else np.nan}
        rows.append(row)
    return pd.DataFrame(rows)


# ======================================================================
#  Main
# ======================================================================
def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--jsonl", required=True)
    ap.add_argument("--gen_root", required=True)
    ap.add_argument("--gt_root", required=True)
    ap.add_argument("--out_dir", default="./mask_population_eval")
    ap.add_argument("--ae_epochs", type=int, default=20)
    ap.add_argument("--ae_size", type=int, default=128)
    ap.add_argument("--latent_dim", type=int, default=64)
    ap.add_argument("--knn_k", type=int, default=5)
    ap.add_argument("--device", default="cuda")
    return ap.parse_args()


def main():
    args = parse_args()
    device = args.device if torch.cuda.is_available() else "cpu"
    out_dir = Path(args.out_dir); out_dir.mkdir(parents=True, exist_ok=True)

    records = collect_masks(args.jsonl, args.gen_root, args.gt_root)
    if not records:
        print("No usable pairs."); return

    # === 1. 训 AE on GT ===
    ae = train_ae([r["gt"] for r in records],
                  latent_dim=args.latent_dim, ae_size=args.ae_size,
                  epochs=args.ae_epochs, device=device)

    # === 2. 编码 gen / gt ===
    feat_gen = encode_masks(ae, [r["gen"] for r in records], args.ae_size, device=device)
    feat_gt  = encode_masks(ae, [r["gt"]  for r in records], args.ae_size, device=device)

    print("\n" + "="*60)
    print("Population-level metrics (overall + per disease)")
    print("="*60)
    diseases = sorted({r["disease"] for r in records})
    pop_rows = []
    # overall
    fid = compute_fid(feat_gen, feat_gt)
    dens, cov = density_coverage(feat_gen, feat_gt, k=args.knn_k)
    pop_rows.append({"disease": "OVERALL", "n": len(records),
                     "Mask_FID": fid, "Density": dens, "Coverage": cov})
    # per disease
    for dis in diseases:
        idx = [i for i, r in enumerate(records) if r["disease"] == dis]
        if len(idx) < 10:
            pop_rows.append({"disease": dis, "n": len(idx),
                             "Mask_FID": np.nan, "Density": np.nan, "Coverage": np.nan})
            continue
        fg, ft = feat_gen[idx], feat_gt[idx]
        fid_d = compute_fid(fg, ft)
        dens_d, cov_d = density_coverage(fg, ft, k=min(args.knn_k, len(idx)-1))
        pop_rows.append({"disease": dis, "n": len(idx),
                         "Mask_FID": fid_d, "Density": dens_d, "Coverage": cov_d})

    pop_df = pd.DataFrame(pop_rows)
    pop_df.to_csv(out_dir / "population_metrics.csv", index=False)
    print(pop_df.round(4).to_string(index=False))

    # === 3. Shape descriptors ===
    print("\n" + "="*60)
    print("Shape descriptors (Wasserstein, lower=better)")
    print("="*60)
    df_gen = shape_descriptors_table(records, "gen")
    df_gt  = shape_descriptors_table(records, "gt")
    df_gen.to_csv(out_dir / "shape_gen.csv", index=False)
    df_gt.to_csv(out_dir / "shape_gt.csv", index=False)
    descriptors = ["area", "aspect_ratio", "circularity",
                   "solidity", "extent", "cx", "cy"]
    w_df = shape_wasserstein(df_gen, df_gt, descriptors)
    w_df.to_csv(out_dir / "shape_wasserstein.csv", index=False)
    print(w_df.round(4).to_string(index=False))

    print("\nPlotting KDE per descriptor ...")
    plot_descriptor_kde(df_gen, df_gt, out_dir / "kde_plots", descriptors)
    print(f"KDE saved -> {out_dir/'kde_plots'}/")

    # === 4. Severity-conditional spread ===
    print("\n" + "="*60)
    print("Severity-conditional spread (between/within > 1 = severity controls form)")
    print("="*60)
    sev_gen = severity_spread(records, ae, args.ae_size, device, side="gen")
    sev_gt  = severity_spread(records, ae, args.ae_size, device, side="gt")
    sev_all = pd.concat([sev_gen, sev_gt], ignore_index=True)
    sev_all.to_csv(out_dir / "severity_spread.csv", index=False)
    print(sev_all.round(3).to_string(index=False))

    print(f"\nAll results saved to {out_dir}/")


if __name__ == "__main__":
    main()

"""
python eval_mask.py \
  --jsonl  /home/jovyan/AURAD_infer/mask/det-test-cp9000/test_prompt_text2layout_single.json \
  --gen_root   /home/jovyan/AURAD_infer/mask/det-test-cp9000 \
  --gt_root  /home/jovyan/AURAD_dataset \
  --out_dir  ./mask_eval \
  --ae_epochs 20 --device cuda
"""