"""
读 evaluate_masks.py 产出的 CSV, 画箱线图.

每个 sample 一个分数 -> 3 个 metric 各画一个箱体, 合并展示.
另外可选: 同时画 3 张独立子图(--mode separate).

用法:
    python plot_boxplot.py --csv mask_scores.csv --out_dir plots
    python plot_boxplot.py --csv mask_scores.csv --out_dir plots --mode separate
"""

import argparse
import os

import matplotlib.pyplot as plt
import pandas as pd


METRICS = ["dice", "iou", "ms_ssim"]
COLORS = ["#4C72B0", "#55A868", "#C44E52"]


def plot_combined(df, out_dir):
    os.makedirs(out_dir, exist_ok=True)
    data = [df[m].dropna().values for m in METRICS]

    fig, ax = plt.subplots(figsize=(7, 6))
    bp = ax.boxplot(
        data,
        labels=[m.upper() for m in METRICS],
        patch_artist=True,
        showmeans=True,
        meanprops=dict(marker="D", markerfacecolor="white",
                       markeredgecolor="black", markersize=6),
    )
    for patch, c in zip(bp["boxes"], COLORS):
        patch.set_facecolor(c)
        patch.set_alpha(0.6)

    ax.set_ylabel("score")
    ax.set_title(f"Per-sample metric distribution (n={len(df)})")
    ax.grid(axis="y", alpha=0.3)
    plt.tight_layout()

    out_path = os.path.join(out_dir, "boxplot_combined.png")
    plt.savefig(out_path, dpi=150)
    plt.close()
    print(f"[saved] {out_path}")


def plot_separate(df, out_dir):
    os.makedirs(out_dir, exist_ok=True)
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    for ax, m, c in zip(axes, METRICS, COLORS):
        vals = df[m].dropna().values
        bp = ax.boxplot(
            [vals],
            labels=[m.upper()],
            patch_artist=True,
            showmeans=True,
            meanprops=dict(marker="D", markerfacecolor="white",
                           markeredgecolor="black", markersize=6),
        )
        bp["boxes"][0].set_facecolor(c)
        bp["boxes"][0].set_alpha(0.6)
        ax.set_title(f"{m.upper()}  (n={len(vals)})\nmean={vals.mean():.3f}  median={pd.Series(vals).median():.3f}")
        ax.grid(axis="y", alpha=0.3)
    plt.tight_layout()
    out_path = os.path.join(out_dir, "boxplot_separate.png")
    plt.savefig(out_path, dpi=150)
    plt.close()
    print(f"[saved] {out_path}")


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--csv", required=True)
    p.add_argument("--out_dir", default="plots")
    p.add_argument("--mode", choices=["combined", "separate", "both"], default="both")
    args = p.parse_args()

    df = pd.read_csv(args.csv)
    df = df[df["status"] == "ok"].copy()
    print(f"[info] {len(df)} ok samples")

    if args.mode in ("combined", "both"):
        plot_combined(df, args.out_dir)
    if args.mode in ("separate", "both"):
        plot_separate(df, args.out_dir)

    if "dataset_fid" in df.columns:
        fid = df["dataset_fid"].dropna().unique()
        if len(fid):
            print(f"[info] dataset FID = {fid[0]}")


if __name__ == "__main__":
    main()


"""
python plot_box.py --csv mask_scores.csv --out_dir figs
"""