#!/usr/bin/env python3
"""Gera boxplots ε × {ASR, SSIM, LPIPS, PSNR} a partir do CSV de um sweep.

Pós-processa o output do `exp_probe_eps_curve_tcc.yaml` (ou qualquer sweep
com múltiplos ε) pra produzir as 4 figuras centrais do §IV.C do paper:
transição "boxplot largo→fino" conforme ε cresce (orientação Maynara 2026-05-05).

Boxplots agrupam por (ε, ataque); cada box agrega ASR/SSIM/LPIPS/PSNR de TODAS
as imagens × modelos rodados naquele (ε, ataque).

Usage:
    # 1. Merge as CSVs do array primeiro:
    python scripts/merge_sweep_results.py \\
        --sweep-dir results/raw/probe_eps_curve_tcc

    # 2. Gerar figuras:
    python scripts/plot_eps_curves_from_sweep_csv.py \\
        --csv results/raw/probe_eps_curve_tcc/merged.csv \\
        --out results/figures/probe_eps_curve_tcc/

Output:
    asr_vs_eps_boxplot.png    — figura central paper (Carlini §5.3)
    ssim_vs_eps_boxplot.png   — descritiva (Sen 2020/Liu 2025)
    lpips_vs_eps_boxplot.png  — descritiva
    psnr_vs_eps_boxplot.png   — descritiva
    summary.md                — Mahmood/TGR comparison + medianas por ε
"""
from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import pandas as pd


def _project_root() -> Path:
    cur = Path(__file__).resolve().parent
    for p in [cur, *cur.parents]:
        if (p / "requirements.txt").exists():
            return p
    raise RuntimeError("project root not found")


PROJECT_ROOT = _project_root()


PALETTE = {
    "FGSM": "#e41a1c", "PGD": "#377eb8", "MIM": "#4daf4a",
    "TGR": "#984ea3", "SAGA": "#ff7f00",
}
ATTACK_ORDER = ["FGSM", "PGD", "MIM", "TGR", "SAGA"]


def boxplot_eps_metric(
    df: pd.DataFrame,
    metric: str,
    ylabel: str,
    title: str,
    out_path: Path,
    ylim: tuple | None = None,
) -> None:
    """Boxplots ε × <metric>, hue=attack."""
    eps_values = sorted(df["epsilon_255"].unique())
    eps_positions = {e: i for i, e in enumerate(eps_values)}
    n_attacks = len(ATTACK_ORDER)
    box_width = 0.8 / n_attacks

    fig, ax = plt.subplots(figsize=(11, 5.5))
    legend_handles = []
    for j, atk in enumerate(ATTACK_ORDER):
        sub = df[df["attack"] == atk]
        if sub.empty:
            continue
        data, positions = [], []
        for e in eps_values:
            vals = sub.loc[sub["epsilon_255"] == e, metric].dropna().values
            if len(vals) == 0:
                continue
            data.append(vals)
            offset = (j - (n_attacks - 1) / 2) * box_width
            positions.append(eps_positions[e] + offset)
        if not data:
            continue
        ax.boxplot(
            data, positions=positions, widths=box_width * 0.85,
            patch_artist=True, showfliers=False,
            medianprops={"color": "black", "linewidth": 1.5},
            boxprops={"facecolor": PALETTE[atk], "alpha": 0.7,
                      "edgecolor": PALETTE[atk]},
            whiskerprops={"color": PALETTE[atk]},
            capprops={"color": PALETTE[atk]},
        )
        legend_handles.append(plt.Rectangle(
            (0, 0), 1, 1, fc=PALETTE[atk], alpha=0.7, label=atk
        ))

    ax.set_xticks(list(eps_positions.values()))
    ax.set_xticklabels([f"{e}/255" for e in eps_values])
    ax.set_xlabel("ε∞ (perturbation budget)")
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    if ylim:
        ax.set_ylim(*ylim)
    ax.grid(axis="y", alpha=0.3)
    ax.legend(handles=legend_handles, loc="best", fontsize=9, ncol=n_attacks)
    fig.tight_layout()
    fig.savefig(out_path, dpi=150)
    plt.close(fig)
    print(f"  ✓ {out_path.name}")


def generate_summary(df: pd.DataFrame, out_path: Path) -> None:
    eps_values = sorted(df["epsilon_255"].unique())
    n_imgs = df["image"].nunique()
    n_models = df["model"].nunique()

    lines = [
        "# Probe ε × métricas — Summary",
        "",
        f"**Imagens**: {n_imgs}",
        f"**Modelos**: {n_models}",
        f"**Ataques**: {sorted(df['attack'].unique())}",
        f"**ε grid**: {eps_values} (×1/255)",
        "",
    ]

    # Por ε: tabela com mediana de ASR / SSIM / LPIPS / PSNR por ataque
    for ref_eps in [4, 8, 16]:
        if ref_eps not in eps_values:
            continue
        sub = df[df["epsilon_255"] == ref_eps]
        lines += [
            f"## Estatística descritiva a ε={ref_eps}/255",
            "",
            "| Attack | ASR | SSIM | LPIPS | PSNR | L∞ |",
            "|---|---|---|---|---|---|",
        ]
        for atk in ATTACK_ORDER:
            sub_atk = sub[sub["attack"] == atk]
            if sub_atk.empty:
                lines.append(f"| {atk} | — | — | — | — | — |")
                continue
            row = (
                f"| {atk} | "
                f"{sub_atk['asr'].mean():.3f} | "
                f"{sub_atk['ssim'].median():.3f} | "
                f"{sub_atk['lpips'].median():.3f} | "
                f"{sub_atk['psnr'].median():.1f} | "
                f"{sub_atk['linf'].median():.4f} |"
            )
            lines.append(row)
        lines.append("")

    # Validação Mahmood/TGR a ε=16/255
    if 16 in eps_values:
        sub16 = df[df["epsilon_255"] == 16]
        lines += [
            "## Validação Mahmood/TGR a ε=16/255",
            "",
            "Esperado (literatura, white-box ImageNet):",
            "  - Mahmood 2021 Tab. 1: ViT-B/16 PGD ≈ 100% ASR, MIM ≈ 100% ASR, FGSM ≈ 76% ASR",
            "  - Zhang 2023 (TGR) Tab. 1: TGR > MIM > PGD > FGSM em transferência (white-box satura)",
            "",
            "Observado (mean ASR sobre modelos × imagens):",
            "",
            "| Attack | Mean ASR | Adv-Acc (1-ASR) |",
            "|---|---|---|",
        ]
        for atk in ATTACK_ORDER:
            sub_atk = sub16[sub16["attack"] == atk]
            if sub_atk.empty:
                lines.append(f"| {atk} | — | — |")
                continue
            asr = sub_atk["asr"].mean()
            lines.append(f"| {atk} | {asr:.3f} | {1-asr:.3f} |")
        lines.append("")

    out_path.write_text("\n".join(lines))
    print(f"  ✓ {out_path.name}")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--csv", type=Path, required=True,
                        help="CSV merged do sweep (output de merge_sweep_results.py)")
    parser.add_argument("--out", type=Path, required=True,
                        help="Diretório de saída pras figuras + summary.md")
    args = parser.parse_args()

    if not args.csv.exists():
        print(f"ERROR: CSV não encontrado: {args.csv}")
        return 1

    args.out.mkdir(parents=True, exist_ok=True)

    print(f"Lendo {args.csv} ...")
    df = pd.read_csv(args.csv)
    df["epsilon_255"] = (df["epsilon"].astype(float) * 255).round().astype(int)
    print(f"  {len(df)} rows, {df['image'].nunique()} imgs × "
          f"{df['model'].nunique()} models × {df['attack'].nunique()} attacks × "
          f"{df['epsilon_255'].nunique()} ε")

    print(f"\nGerando figuras em {args.out}/ ...")
    boxplot_eps_metric(
        df, "asr", "ASR (per image, 0=fail, 1=success)",
        "Curva ε × ASR — distribuição por ataque",
        args.out / "asr_vs_eps_boxplot.png", ylim=(-0.05, 1.05),
    )
    boxplot_eps_metric(
        df, "ssim", "SSIM (higher = more similar)",
        "Curva ε × SSIM — descritiva",
        args.out / "ssim_vs_eps_boxplot.png", ylim=(0.4, 1.02),
    )
    boxplot_eps_metric(
        df, "lpips", "LPIPS (lower = more similar)",
        "Curva ε × LPIPS — descritiva",
        args.out / "lpips_vs_eps_boxplot.png",
    )
    boxplot_eps_metric(
        df, "psnr", "PSNR (dB, higher = more similar)",
        "Curva ε × PSNR — descritiva",
        args.out / "psnr_vs_eps_boxplot.png",
    )

    print(f"\nGerando summary.md ...")
    generate_summary(df, args.out / "summary.md")

    print(f"\n✓ Done. Output: {args.out}")
    return 0


if __name__ == "__main__":
    import sys
    sys.exit(main())