"""
逐样本指标 CSV + Gen vs GT 带统计显著性的箱线图
================================================

输入: 与 eval_mask_population.py 一致
  --jsonl, --gen_root, --gt_root

输出:
  per_sample_shape.csv   每个 mask 的形态学描述子 (gen + gt 各一行)
  boxplot_shape.png      2x2 箱线图: area / aspect_ratio / circularity / solidity
                         每个子图按疾病分面, gen vs GT 两个箱, 带 Mann-Whitney U 显著性星号

用法:
  python plot_shape_boxplot.py \\
      --jsonl /path/test.json \\
      --gen_root /home/jovyan/AURAD_infer/mask/det-test-cp9000 \\
      --gt_root  /home/jovyan/AURAD_dataset \\
      --out_dir  ./boxplot_out
"""
import argparse
import json
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from PIL import Image
from scipy.ndimage import binary_erosion
from scipy.spatial import ConvexHull
from scipy.stats import mannwhitneyu


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


def shape_descriptors(mask):
    if mask.sum() == 0:
        return {k: np.nan for k in
                ["area", "aspect_ratio", "circularity", "solidity"]}
    ys, xs = np.where(mask > 0)
    area = int(len(ys))
    h = ys.max() - ys.min() + 1
    w = xs.max() - xs.min() + 1
    aspect = h / w if w > 0 else np.nan

    eroded = binary_erosion(mask)
    perimeter = int((mask & ~eroded).sum())
    circ = 4 * np.pi * area / (perimeter ** 2) if perimeter > 0 else np.nan

    try:
        pts = np.column_stack([xs, ys])
        if len(pts) >= 3:
            hull = ConvexHull(pts)
            solidity = area / hull.volume
        else:
            solidity = 1.0
    except Exception:
        solidity = np.nan

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


# ======================================================================
def collect_descriptors(jsonl, gen_root, gt_root):
    gen_root, gt_root = Path(gen_root), Path(gt_root)
    items = [json.loads(l) for l in open(jsonl) if l.strip()]
    rows = []
    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
        img_id = Path(it["file_name"]).parent.name
        disease = it["attn_list"][0][0]

        for side, path in [("gen", gen_p), ("gt", gt_p)]:
            d = shape_descriptors(load_mask(path))
            d.update({"id": img_id, "disease": disease, "source": side})
            rows.append(d)
    print(f"Computed descriptors for {len(rows)//2} mask pairs (skipped {n_skip})")
    return pd.DataFrame(rows)


# ======================================================================
def sig_label(p):
    if np.isnan(p):
        return ""
    if p < 1e-4:
        return "****"
    if p < 1e-3:
        return "***"
    if p < 1e-2:
        return "**"
    if p < 5e-2:
        return "*"
    return "ns"


def plot_boxplot(df, out_path):
    metrics = ["area", "aspect_ratio", "circularity", "solidity"]
    titles = {
        "area": "Area (log10)",
        "aspect_ratio": "Aspect Ratio (h/w)",
        "circularity": "Circularity",
        "solidity": "Solidity",
    }
    diseases = sorted(df["disease"].unique())

    # 4 行 1 列, 每行一个指标横向展开所有疾病
    fig, axes = plt.subplots(4, 1, figsize=(max(len(diseases)*1.0 + 2, 12), 16))

    for ax, metric in zip(axes, metrics):
        gen_data, gt_data, labels, p_values = [], [], [], []
        for dis in diseases:
            g = df[(df.disease == dis) & (df.source == "gen")][metric].dropna().values
            t = df[(df.disease == dis) & (df.source == "gt")][metric].dropna().values
            if metric == "area":
                g = np.log10(g + 1)
                t = np.log10(t + 1)
            gen_data.append(g)
            gt_data.append(t)
            labels.append(f"{dis}\n(n={len(t)})")
            if len(g) >= 3 and len(t) >= 3:
                _, p = mannwhitneyu(g, t, alternative="two-sided")
            else:
                p = np.nan
            p_values.append(p)

        x = np.arange(len(diseases))
        width = 0.35
        bp_gen = ax.boxplot(gen_data, positions=x - width/2, widths=width*0.9,
                            patch_artist=True, showfliers=False,
                            medianprops=dict(color="black", linewidth=1.5))
        bp_gt = ax.boxplot(gt_data, positions=x + width/2, widths=width*0.9,
                           patch_artist=True, showfliers=False,
                           medianprops=dict(color="black", linewidth=1.5))
        for b in bp_gen["boxes"]:
            b.set_facecolor("#E89A6B"); b.set_edgecolor("#a66033"); b.set_alpha(0.85)
        for b in bp_gt["boxes"]:
            b.set_facecolor("#4DBBA1"); b.set_edgecolor("#2c7a66"); b.set_alpha(0.85)

        # 显著性星号
        for i, p in enumerate(p_values):
            label = sig_label(p)
            if not label:
                continue
            all_vals = np.concatenate([gen_data[i], gt_data[i]]) if \
                (len(gen_data[i]) and len(gt_data[i])) else np.array([0])
            ytop = np.percentile(all_vals, 95) if len(all_vals) > 5 else (
                all_vals.max() if len(all_vals) else 0)
            yrange = ax.get_ylim()
            yoff = (yrange[1] - yrange[0]) * 0.02 if yrange[1] > yrange[0] else 0.02
            color = "#888" if label == "ns" else "black"
            fs = 9 if label == "ns" else 12
            ax.text(i, ytop + yoff, label, ha="center", va="bottom",
                    fontsize=fs, color=color, weight="bold")

        ax.set_xticks(x)
        ax.set_xticklabels(labels, rotation=30, ha="right", fontsize=9)
        ax.set_ylabel(titles[metric], fontsize=12, weight="bold")
        ax.grid(axis="y", linestyle="--", alpha=0.4)
        ax.set_axisbelow(True)

    from matplotlib.patches import Patch
    handles = [
        Patch(facecolor="#E89A6B", edgecolor="#a66033", label="Generated"),
        Patch(facecolor="#4DBBA1", edgecolor="#2c7a66", label="Real (GT)"),
    ]
    fig.legend(handles=handles, loc="upper center", bbox_to_anchor=(0.5, 0.995),
               ncol=2, fontsize=14, frameon=False)
    fig.text(0.5, 0.005,
             "Mann-Whitney U:  **** p<1e-4   *** p<1e-3   ** p<1e-2   * p<5e-2   ns p>=5e-2",
             ha="center", fontsize=10, color="#555")
    fig.tight_layout(rect=[0, 0.015, 1, 0.97])
    fig.savefig(out_path, dpi=180, bbox_inches="tight")
    plt.close(fig)
    print(f"Boxplot saved: {out_path}")


# ======================================================================
def main():
    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="./boxplot_out")
    args = ap.parse_args()

    out_dir = Path(args.out_dir); out_dir.mkdir(parents=True, exist_ok=True)

    df = collect_descriptors(args.jsonl, args.gen_root, args.gt_root)
    csv_path = out_dir / "per_sample_shape.csv"
    df.to_csv(csv_path, index=False)
    print(f"Saved per-sample CSV: {csv_path}  ({len(df)} rows)")

    plot_boxplot(df, out_dir / "boxplot_shape.png")


if __name__ == "__main__":
    main()


"""
python plot_boxplot.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 
"""