#!/usr/bin/env python3
"""Sweep B-probe: curva ε × {ASR, SSIM, LPIPS, PSNR} + PGD T-plateau check.

Updated 2026-05-08:
  - Dataset: data/in1k_hybrid_1k/ (1 imagem por classe IN-1k, com flag
    has_mask via metadata.json). Amostragem estratificada via metadata.
  - Modelos: 4 do TCC — **S/16, S/32, B/16, B/32** (drop ViT-L). Passar via
    --models. Tabela 2×2 (size × patch).
  - ε grid: 7 pontos {2, 4, 6, 8, 10, 12, 16}/255 (era 5).
  - Plots: boxplots ε × métrica (era linhas mean±IC) — orientação Maynara
    2026-05-05: transição "boxplot largo→fino" é achado científico.
  - Plots novos: psnr_vs_eps_boxplot.png + asr_vs_eps_boxplot.png.

Satisfaz duas recomendações de Carlini et al. (2019):
  - §4.8: verifica que PGD converge por T=10 (plot ASR vs T, espera plateau)
  - §5.3: gera curva accuracy-vs-perturbation (figura central do paper)

Coleta também stats descritivos de qualidade (SSIM/LPIPS/PSNR/L∞) em cada ε
pra reportar como evidência na Discussion (com caveat Sen 2020 + Liu 2025).

Statistical basis for N=100
───────────────────────────
- ASR binomial CI at N=100, p=0.5: ±0.098 — sufficient to distinguish 0% vs 30%
- SSIM mean CI at N=100, σ=0.05: ±0.0098 — sufficient for descriptive curves
- Cost ~15-20 min local (4 modelos × 5 ataques × 7 ε + PGD T-sweep)

Attack hyperparameters (paper-exact, see protocolo-avaliacao-imagenet.md)
──────────────────────────────────────────────────────────────────────────
  FGSM  T=1  α=ε            (Goodfellow et al., ICLR 2015 — single step)
  PGD   T=10 α=ε/T  random_start=False  (Madry 2018, with random_start=False
                                          for reproducibility of attention maps)
  MIM   T=10 α=ε/T  μ=1.0   (Dong et al., CVPR 2018 — Sec. 4 ImageNet)
  TGR   T=10 α=ε/T  μ=1.0  (Zhang et al., CVPR 2023 — k=1, γ_attn=0.25,
                              γ_QKV=0.75, γ_MLP=0.25 (paper Tab. 3))
  SAGA  T=10 α=ε/T          (Mahmood et al., ICCV 2021 — white-box ensemble)

Note on MIM alpha: Dong 2018 Sec. 4 ImageNet uses α = ε/T = 1.6/255 at ε=16/255
  (proportional, NOT fixed at 2/255 as previously thought — corrected 2026-05-02).
Note on TGR gamma_mlp: paper Tab. 3 ablation → s_MLP=0.25 (code default is 0.5;
  overridden here for paper-exact reproduction).

Usage
─────
  # TCC scope: 4 modelos S/B × 5 ataques × 7 ε no dataset híbrido 1K
  python experiments/run_imperceptibility_probe.py \\
      --models "ViT-S/16,ViT-S/32,ViT-B/16,ViT-B/32"

  # Custom images directory
  python experiments/run_imperceptibility_probe.py --images-dir /path/to/imgs

  # Skip PGD T-plateau check
  python experiments/run_imperceptibility_probe.py --no-t-sweep

  # Regenerate plots from existing CSV (no re-run)
  python experiments/run_imperceptibility_probe.py --plot-only

  # Resume an interrupted run
  python experiments/run_imperceptibility_probe.py --resume

Outputs (under results/imperceptibility_probe/)
────────────────────────────────────────────────
  probe_results.csv          — raw rows: ... + has_mask (1=Guillaumin GT, 0=outros)
  asr_vs_eps_boxplot.png     — Carlini §5.3 figura central — boxplots ε × ASR
  ssim_vs_eps_boxplot.png    — boxplot ε × SSIM (descritiva, Sen 2020/Liu 2025)
  lpips_vs_eps_boxplot.png   — boxplot ε × LPIPS
  psnr_vs_eps_boxplot.png    — boxplot ε × PSNR
  pgd_t_plateau.png          — Carlini §4.8 sanity — ASR(T) para PGD a ε=8/255
  summary.md                 — T-plateau verdict + descriptive stats + Mahmood/TGR
                               comparison at ε=16/255
"""

from __future__ import annotations

import argparse
import csv
import json
import random
import sys
import time
from pathlib import Path

# ─── project root ────────────────────────────────────────────────────────────

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

ROOT = _project_root()
sys.path.insert(0, str(ROOT))

import numpy as np
import torch
import yaml
from tqdm import tqdm

from utils.attacks import FGSM, MIFGSM, PGDIterations, SAGA, TGR
from utils.metrics import compute_all_image_metrics, compute_all_attack_metrics
from utils.model_loader import load_model_and_labels
from utils.preprocessing import get_default_transform, preprocess_image
from utils.seed import set_seed

# ─── constants ───────────────────────────────────────────────────────────────

#: ε grid (L∞, in [0,1] space). 7 points covering the range from imperceptible
#: (2/255, Goodfellow 2015 panda demo) to perceptible (16/255, Mahmood/TGR).
#: Inclui {6, 10}/255 para visualizar transição "boxplot largo → fino" entre
#: regimes (Maynara 2026-05-05). Usado para gerar curvas ε × {ASR, SSIM, LPIPS,
#: PSNR} centrais ao paper (Carlini §5.3 + decisão protocolo 2026-05-05).
EPSILON_GRID = [e / 255 for e in [2, 4, 6, 8, 10, 12, 16]]

#: T-sweep for PGD plateau check (Carlini §4.8). Run at fixed ε=8/255 (Sweep B
#: principal). Plateau verdict: ASR(T=40) − ASR(T=20) < 1pp.
PGD_T_PLATEAU_EPS = 8 / 255
PGD_T_GRID = [5, 10, 20, 40]

#: Default number of images. N=200 gives ASR CI ±0.069 at p=0.5 (apertado pra
#: regime distinction) e SSIM CI ±0.0069 (descritiva). Sampleadas
#: estratificadamente do data/in1k_hybrid_1k/ via metadata.json (todos os 95
#: com Guillaumin GT mask + 105 sem mask sorteados).
#: Cost: ~30-40 min local em MPS / ~15-20 min A100 (4 models × 5 attacks × 7 ε
#: × 200 imgs = 28k combos main + 4×4×200 = 3.2k T-sweep PGD).
DEFAULT_N_IMAGES = 200

#: Reference epsilon for Mahmood/TGR comparison (paper validation).
MAHMOOD_TGR_EPS = 16 / 255

OUTPUT_DIR = ROOT / "results" / "imperceptibility_probe"

CSV_COLUMNS = [
    "model", "attack", "epsilon", "seed", "image", "has_mask",
    "orig_pred", "orig_conf", "adv_pred", "adv_conf",
    "asr", "confidence_drop", "topk_drop",
    "linf", "psnr", "ssim", "lpips", "modified_pixels",
    "time_seconds",
]

# ─── paper-exact attack hyperparameters ──────────────────────────────────────

#: All non-ε hyperparameters fixed to paper originals. Only ε is swept.
PAPER_ATTACKS: list[dict] = [
    {
        "name": "FGSM",
        "label": "FGSM (Goodfellow 2015)",
        "ref": "ICLR 2015",
        # T=1, α=ε — single-step sign gradient (no additional params needed)
    },
    {
        "name": "PGD",
        "label": "PGD (Madry 2018)",
        "ref": "ICLR 2018",
        "steps": 10,
        # α = ε/T — Madry formulation; for T=10 → alpha_ratio=0.1
        "alpha_ratio": 0.1,
        "random_start": False,
    },
    {
        "name": "MIM",
        "label": "MIM (Dong 2018)",
        "ref": "CVPR 2018",
        "steps": 10,
        # CORRECTED 2026-05-02: Dong 2018 Sec. 4 ImageNet uses α = ε/T,
        # NOT α=2/255 fixed as previously documented. With ε=16, T=10 → α=1.6.
        # Verification: 1.6 = 16/10 = ε/T (proportional, not fixed).
        # alpha_ratio=0.1 reproduces paper-exact α at any ε.
        "alpha_ratio": 0.1,
        "decay": 1.0,  # μ = 1.0 (momentum decay factor)
    },
    {
        "name": "TGR",
        "label": "TGR (Zhang 2023)",
        "ref": "CVPR 2023",
        "steps": 10,
        # α = ε/T internally (eps_step = eps/steps in TGR.__init__)
        # At paper's ε=16/255, T=10 → α=1.6/255 ✓
        "decay": 1.0,         # μ = 1.0
        "k": 1,               # paper Table 3: k=1
        "gamma_attn": 0.25,   # s_attn = 0.25 (paper Table 3)
        "gamma_qkv": 0.75,    # s_QKV  = 0.75 (paper Table 3)
        "gamma_mlp": 0.25,    # s_MLP  = 0.25 (paper Table 3; code default=0.5, OVERRIDDEN)
    },
    {
        "name": "SAGA",
        "label": "SAGA (Mahmood 2021)",
        "ref": "ICCV 2021",
        "steps": 10,
        # α = ε/T internally (eps_step = eps/steps in SAGA.__init__)
    },
]


def build_attack(model: torch.nn.Module, attack_def: dict, eps: float):
    """Instantiate an attack with exact paper hyperparameters.

    Only ε is variable. All other params come from PAPER_ATTACKS definitions.
    """
    name = attack_def["name"]

    if name == "FGSM":
        return FGSM(model, eps=eps, collect_images=False)

    elif name == "PGD":
        alpha = eps * attack_def["alpha_ratio"]  # ε/T
        return PGDIterations(
            model,
            eps=eps,
            alpha=alpha,
            steps=attack_def["steps"],
            random_start=attack_def["random_start"],
            collect_images=False,
        )

    elif name == "MIM":
        alpha = eps * attack_def["alpha_ratio"]  # α = ε/T = 0.1·ε (Dong 2018 Sec. 4)
        return MIFGSM(
            model,
            eps=eps,
            alpha=alpha,
            steps=attack_def["steps"],
            decay=attack_def["decay"],
            collect_images=False,
        )

    elif name == "TGR":
        return TGR(
            model,
            eps=eps,
            steps=attack_def["steps"],
            decay=attack_def["decay"],
            k=attack_def["k"],
            gamma_attn=attack_def["gamma_attn"],
            gamma_qkv=attack_def["gamma_qkv"],
            gamma_mlp=attack_def["gamma_mlp"],  # 0.25, not code default 0.5
            collect_images=False,
        )

    elif name == "SAGA":
        return SAGA(model, eps=eps, steps=attack_def["steps"], collect_images=False)

    else:
        raise ValueError(f"Unknown attack: {name}")


# ─── config + model loading ───────────────────────────────────────────────────

def load_config(config_path: str | None) -> dict:
    default = ROOT / "configs" / "default.yaml"
    with open(default) as f:
        cfg = yaml.safe_load(f)
    if config_path and Path(config_path).exists():
        with open(config_path) as f:
            overrides = yaml.safe_load(f) or {}
        for k, v in overrides.items():
            if isinstance(v, dict) and k in cfg and isinstance(cfg[k], dict):
                cfg[k].update(v)
            else:
                cfg[k] = v
    return cfg


def get_device(cfg: dict) -> torch.device:
    dev = cfg.get("device", "auto")
    if dev == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(dev)


# ─── single-run logic ─────────────────────────────────────────────────────────

def run_single(
    model, vit_config, attack_def: dict, eps: float,
    img_path: Path, device: torch.device, seed: int,
) -> dict:
    set_seed(seed)

    transform = get_default_transform(img_size=vit_config.img_size)
    img_tensor = preprocess_image(str(img_path), transform=transform).to(device)

    model.eval()
    with torch.no_grad():
        out = model(img_tensor)
        logits = (out[0] if isinstance(out, tuple) else out)[0]
        probs = torch.nn.functional.softmax(logits, dim=0)
        orig_pred = probs.argmax().item()
        orig_conf = probs[orig_pred].item()

    attack = build_attack(model, attack_def, eps)
    label = torch.tensor([orig_pred], device=device)

    t0 = time.time()
    adv_tensor, _ = attack(img_tensor, label)
    elapsed = time.time() - t0

    with torch.no_grad():
        out_adv = model(adv_tensor)
        logits_adv = (out_adv[0] if isinstance(out_adv, tuple) else out_adv)[0]
        probs_adv = torch.nn.functional.softmax(logits_adv, dim=0)
        adv_pred = probs_adv.argmax().item()
        adv_conf = probs_adv[adv_pred].item()

    img_m = compute_all_image_metrics(img_tensor, adv_tensor, use_lpips=True)
    atk_m = compute_all_attack_metrics(probs, probs_adv, orig_pred, adv_pred)

    return {
        "orig_pred": orig_pred,
        "orig_conf": orig_conf,
        "adv_pred": adv_pred,
        "adv_conf": adv_conf,
        "elapsed": elapsed,
        **img_m,
        **atk_m,
    }


# ─── CSV helpers ──────────────────────────────────────────────────────────────

def _fmt(v) -> str:
    if isinstance(v, float):
        return "" if (v != v) else f"{v:.6f}"  # NaN → empty
    return str(v)


def completed_keys(csv_path: Path) -> set:
    if not csv_path.exists():
        return set()
    keys: set = set()
    with open(csv_path) as f:
        for row in csv.DictReader(f):
            keys.add((row["model"], row["attack"], row["epsilon"], row["seed"], row["image"]))
    return keys


# ─── plotting ─────────────────────────────────────────────────────────────────

def generate_plots(csv_path: Path, out_dir: Path) -> None:
    """Read probe CSV and produce boxplots ε × {ASR, SSIM, LPIPS, PSNR} + PGD T-plateau.

    Boxplots ao invés de linhas mean±IC (orientação Maynara 2026-05-05): a
    transição "boxplot largo → fino" conforme ε cresce é achado científico em si
    — em ε baixo, atacabilidade depende da imagem (heterogeneidade); em ε alto,
    satura universalmente. ε=8/255 cai na transição (sweet spot).
    """
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        import pandas as pd
    except ImportError:
        print("[plots] matplotlib/pandas not available — skipping plots.")
        return

    out_dir.mkdir(parents=True, exist_ok=True)
    df = pd.read_csv(csv_path)
    df["epsilon_255"] = (df["epsilon"].astype(float) * 255).round(0).astype(int)

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

    main_df = df[df["t_sweep"].isna() | (df["t_sweep"] == "")] if "t_sweep" in df.columns else df

    # ε ticks (ordenados, todos os valores presentes)
    eps_values = sorted(main_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  # largura individual de cada box dentro do grupo ε

    def _boxplot_eps_metric(metric: str, ylabel: str, title: str, fname: str,
                            ylim: tuple | None = None,
                            higher_is_better: bool = True) -> None:
        fig, ax = plt.subplots(figsize=(11, 5.5))
        legend_handles = []
        for j, atk in enumerate(attack_order):
            sub = main_df[main_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)
                # posição: centro do grupo ε + offset por ataque
                offset = (j - (n_attacks - 1) / 2) * box_width
                positions.append(eps_positions[e] + offset)
            if not data:
                continue
            bp = 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]},
            )
            # Manter handle pra legenda
            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_dir / fname, dpi=150)
        plt.close(fig)

    # ── ASR vs ε (figura central — Carlini §5.3) ──────────────────────────────
    _boxplot_eps_metric(
        "asr", ylabel="ASR (per image, 0=fail, 1=success)",
        title="Curva ε × ASR — distribuição por ataque (boxplots)",
        fname="asr_vs_eps_boxplot.png", ylim=(-0.05, 1.05),
    )

    # ── SSIM vs ε (descritiva — Sen 2020/Liu 2025) ────────────────────────────
    _boxplot_eps_metric(
        "ssim", ylabel="SSIM (higher = more similar)",
        title="Curva ε × SSIM — descritiva (Wang 2004 não fornece threshold)",
        fname="ssim_vs_eps_boxplot.png", ylim=(0.4, 1.02),
    )

    # ── LPIPS vs ε ────────────────────────────────────────────────────────────
    _boxplot_eps_metric(
        "lpips", ylabel="LPIPS (lower = more similar)",
        title="Curva ε × LPIPS — descritiva",
        fname="lpips_vs_eps_boxplot.png",
    )

    # ── PSNR vs ε ─────────────────────────────────────────────────────────────
    _boxplot_eps_metric(
        "psnr", ylabel="PSNR (dB, higher = more similar)",
        title="Curva ε × PSNR — descritiva",
        fname="psnr_vs_eps_boxplot.png",
    )

    # ── PGD T-plateau (Carlini §4.8) ──
    if "t_sweep" in df.columns:
        t_df = df[df["t_sweep"].notna() & (df["t_sweep"] != "")]
        if not t_df.empty:
            t_df = t_df.copy()
            t_df["steps"] = t_df["t_sweep"].astype(int)
            fig, ax = plt.subplots(figsize=(6, 4))
            grp = t_df.groupby("steps")["asr"].mean().reset_index()
            ax.plot(grp["steps"], grp["asr"], marker="o", color=colors["PGD"],
                    linewidth=2, markersize=7)
            ax.set_xlabel("PGD iterations (T)")
            ax.set_ylabel("ASR (mean over models)")
            ax.set_title(f"Carlini §4.8 — PGD T-plateau at ε={PGD_T_PLATEAU_EPS*255:.0f}/255")
            ax.set_xticks(PGD_T_GRID)
            ax.grid(alpha=0.3)
            fig.tight_layout()
            fig.savefig(out_dir / "pgd_t_plateau.png", dpi=150)
            plt.close(fig)

    print(f"[plots] saved to {out_dir}/")


def generate_summary(csv_path: Path, out_dir: Path) -> None:
    """Print and save a Markdown summary with T-plateau verdict, descriptive
    stats at ε=8/255 and ε=16/255, and Mahmood/TGR comparison.
    """
    try:
        import pandas as pd
    except ImportError:
        print("[summary] pandas not available — skipping summary.")
        return

    df = pd.read_csv(csv_path)
    df["epsilon_255"] = (df["epsilon"].astype(float) * 255).round(1)
    main_df = df[df["t_sweep"].isna() | (df["t_sweep"] == "")] if "t_sweep" in df.columns else df

    lines = [
        "# Sweep B-probe — Summary",
        "",
        f"**N images**: {main_df['image'].nunique()}  ",
        f"**Models**: {main_df['model'].nunique()}  ",
        f"**Attacks**: {sorted(main_df['attack'].unique())}  ",
        f"**ε grid**: {sorted(main_df['epsilon_255'].unique())} (×1/255)  ",
        "",
        "## 1. PGD T-plateau verdict (Carlini 2019 §4.8)",
        "",
    ]

    # T-plateau check
    if "t_sweep" in df.columns:
        t_df = df[df["t_sweep"].notna() & (df["t_sweep"] != "")].copy()
        if not t_df.empty:
            t_df["steps"] = t_df["t_sweep"].astype(int)
            grp = t_df.groupby("steps")["asr"].mean().reset_index().sort_values("steps")
            lines.append("| T  | ASR (mean over models) |")
            lines.append("|---|---|")
            for _, r in grp.iterrows():
                lines.append(f"| {int(r['steps']):2d} | {r['asr']:.4f} |")

            asr_by_t = dict(zip(grp["steps"], grp["asr"]))
            if 20 in asr_by_t and 40 in asr_by_t:
                delta = abs(asr_by_t[40] - asr_by_t[20])
                verdict = "✅ PASS" if delta < 0.01 else "⚠️  CHECK"
                lines.append("")
                lines.append(f"|ASR(T=40) − ASR(T=20)| = {delta:.4f} → {verdict} "
                             f"(threshold: <0.01 = <1pp)")
                if delta < 0.01:
                    lines.append("Plateau confirmed: T=10 is sufficient for the main sweep.")
                else:
                    lines.append("Plateau NOT confirmed: consider raising T to 20.")
        else:
            lines.append("(No T-sweep data — use --no-t-sweep flag was set or sweep skipped)")
    else:
        lines.append("(No T-sweep column in CSV)")

    # Descriptive stats at ε=8/255 and ε=16/255
    for ref_eps in [8.0, 16.0]:
        sub = main_df[main_df["epsilon_255"] == ref_eps]
        if sub.empty:
            continue

        lines += [
            "",
            f"## 2. Descriptive stats at ε={int(ref_eps)}/255 (mean over images × models)",
            "",
            "| Attack | ASR  | SSIM  | LPIPS | PSNR  | L∞    |",
            "|---|---|---|---|---|---|",
        ]
        for atk in ["FGSM", "PGD", "MIM", "TGR", "SAGA"]:
            atk_sub = sub[sub["attack"] == atk]
            if atk_sub.empty:
                lines.append(f"| {atk} | — | — | — | — | — |")
                continue
            row = (
                f"| {atk:<5} | "
                f"{atk_sub['asr'].mean():.3f} | "
                f"{atk_sub['ssim'].mean():.3f} | "
                f"{atk_sub['lpips'].mean():.3f} | "
                f"{atk_sub['psnr'].mean():5.2f} | "
                f"{atk_sub['linf'].mean():.4f} |"
            )
            lines.append(row)

    # Mahmood/TGR validation — only at ε=16/255
    sub16 = main_df[main_df["epsilon_255"] == 16.0]
    if not sub16.empty:
        lines += [
            "",
            "## 3. Mahmood/TGR comparison at ε=16/255 (implementation validation)",
            "",
            "Expected (from published baselines):",
            "  - Mahmood 2021 Tab. 1 (ImageNet, ε=0.062): ViT-B/16 PGD ≈ 0% Acc, MIM ≈ 0% Acc, FGSM ≈ 24% Acc",
            "  - Zhang 2023 (TGR) Tab. 1: TGR > MIM > PGD > FGSM in transfer; for white-box should be saturated",
            "",
            "Observed (mean Adv-Acc = 1 - ASR over models):",
            "",
            "| Attack | Adv-Acc | ASR | Note |",
            "|---|---|---|---|",
        ]
        for atk in ["FGSM", "PGD", "MIM", "TGR", "SAGA"]:
            atk_sub = sub16[sub16["attack"] == atk]
            if atk_sub.empty:
                lines.append(f"| {atk} | — | — | no data |")
                continue
            asr = atk_sub["asr"].mean()
            adv_acc = 1.0 - asr
            note = "saturated (expected)" if asr > 0.95 else "not saturated"
            lines.append(f"| {atk} | {adv_acc:.3f} | {asr:.3f} | {note} |")

    lines += [
        "",
        "## References",
        "- Carlini et al. (2019) — On Evaluating Adversarial Robustness. arXiv 1902.06705 §4.8 §5.3",
        "- Goodfellow et al. (2015) — FGSM. ICLR 2015.",
        "- Madry et al. (2018) — PGD. ICLR 2018.",
        "- Dong et al. (2018) — MIM. CVPR 2018. (α=ε/T per Sec. 4)",
        "- Zhang et al. (2023) — TGR. CVPR 2023. (γ_MLP=0.25 per Tab. 3)",
        "- Mahmood et al. (2021) — SAGA. ICCV 2021.",
        "- Sen et al. (2020) — Imperceptibility measures far from human perception. GameSec 2020.",
        "- Liu et al. (2025) — Stealthiness Assessment of Adversarial Perturbation. IEEE TIFS.",
    ]

    md = "\n".join(lines)
    summary_path = out_dir / "summary.md"
    summary_path.write_text(md)
    print(md)
    print(f"\n[summary] saved to {summary_path}")


# ─── main ─────────────────────────────────────────────────────────────────────

def main() -> None:
    parser = argparse.ArgumentParser(description="Sweep B-probe: imperceptibility calibration")
    parser.add_argument("--config", default=None, help="Experiment YAML (for model paths)")
    parser.add_argument("--models", default=None,
                        help='Comma-separated model name prefixes to filter (e.g. '
                             '"ViT-S/16,ViT-S/32,ViT-B/16,ViT-B/32" para TCC scope). '
                             "Default: usa todos do config (default.yaml = 6 modelos).")
    parser.add_argument("--images-dir", default=None,
                        help="Directory with probe images (default: data/in1k_hybrid_1k/)")
    parser.add_argument("--n-images", type=int, default=DEFAULT_N_IMAGES,
                        help=f"Number of images to use (default: {DEFAULT_N_IMAGES}, min: 100)")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device", default=None)
    parser.add_argument("--output-dir", default=str(OUTPUT_DIR))
    parser.add_argument("--no-t-sweep", action="store_true",
                        help="Skip PGD T-plateau check (Carlini §4.8)")
    parser.add_argument("--plot-only", action="store_true",
                        help="Skip sweep, regenerate plots from existing CSV")
    parser.add_argument("--resume", action="store_true", help="Skip already-computed rows")
    parser.add_argument("--dry-run", action="store_true", help="List combinations without running")
    args = parser.parse_args()

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    csv_path = out_dir / "probe_results.csv"

    if args.plot_only:
        if not csv_path.exists():
            print(f"ERROR: {csv_path} not found. Run without --plot-only first.")
            sys.exit(1)
        generate_plots(csv_path, out_dir)
        generate_summary(csv_path, out_dir)
        return

    # Load config + resolve paths
    cfg = load_config(args.config)
    device = torch.device(args.device) if args.device else get_device(cfg)
    images_dir = Path(args.images_dir) if args.images_dir else ROOT / "data" / "in1k_hybrid_1k"

    if not images_dir.exists():
        print(f"ERROR: images directory not found: {images_dir}")
        print("  → run `python scripts/build_in1k_hybrid_1k.py --seed 42` first")
        sys.exit(1)

    # Filtrar modelos por nome se --models foi passado
    if args.models:
        wanted_names = {m.strip() for m in args.models.split(",")}
        cfg["models"] = [m for m in cfg["models"]
                         if m["name"].split(" ·")[0].strip() in wanted_names
                         or m["name"].strip() in wanted_names]
        if not cfg["models"]:
            print(f"ERROR: no models matched filter --models={args.models}")
            print(f"  Available: {[m['name'].split(' ·')[0] for m in load_config(args.config)['models']]}")
            sys.exit(1)

    # Amostragem estratificada via metadata.json (prioriza has_mask=True)
    metadata_path = images_dir / "metadata.json"
    has_mask_lookup: dict[str, bool] = {}
    if metadata_path.exists():
        meta = json.loads(metadata_path.read_text())
        samples = meta.get("samples") or []
        if samples:
            with_mask = [s["filename"] for s in samples if s.get("has_mask")]
            without_mask = [s["filename"] for s in samples if not s.get("has_mask")]
            has_mask_lookup = {s["filename"]: bool(s.get("has_mask"))
                               for s in samples}
            rng = random.Random(args.seed)
            n = min(args.n_images, len(samples))
            if n <= len(with_mask):
                chosen_names = rng.sample(with_mask, n)
            else:
                chosen_names = with_mask + rng.sample(
                    without_mask, n - len(with_mask)
                )
            chosen_names = sorted(chosen_names)
            images = [images_dir / nm for nm in chosen_names if (images_dir / nm).exists()]
            print(f"Stratified sample: {n} imgs ({sum(1 for p in images if has_mask_lookup.get(p.name)):d} com mask)")
        else:
            images = []
    else:
        # Fallback: sem metadata, iterar arquivos lexicograficamente
        exts = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}
        all_images = sorted(p for p in images_dir.iterdir() if p.suffix.lower() in exts)
        n = min(args.n_images, len(all_images))
        images = all_images[:n]
        print(f"WARN: metadata.json não encontrado em {images_dir}; usando primeiras {n} imagens lexicograficamente")

    if len(images) < 100:
        print(f"WARNING: only {len(images)} images available (recommended ≥100 for reliable CI)")
    n = len(images)

    # Build combinations: (model_cfg, attack_def, eps, image)
    # T-sweep for PGD adds extra combinations at fixed ε
    combos: list[tuple] = []
    for model_cfg in cfg["models"]:
        for attack_def in PAPER_ATTACKS:
            for eps in EPSILON_GRID:
                for img in images:
                    combos.append((model_cfg, attack_def, eps, img, None))  # None = not t-sweep

    if not args.no_t_sweep:
        pgd_def = next(a for a in PAPER_ATTACKS if a["name"] == "PGD")
        for model_cfg in cfg["models"]:
            for t in PGD_T_GRID:
                pgd_t = dict(pgd_def, steps=t, alpha_ratio=1 / t)
                for img in images:
                    combos.append((model_cfg, pgd_t, PGD_T_PLATEAU_EPS, img, t))

    n_models = len(cfg["models"])
    n_eps = len(EPSILON_GRID)
    n_attacks = len(PAPER_ATTACKS)
    t_extra = 0 if args.no_t_sweep else n_models * len(PGD_T_GRID) * n
    print(
        f"\nSweep B-probe configuration"
        f"\n  Device:        {device}"
        f"\n  Images:        {n} (from {images_dir})"
        f"\n  Models:        {n_models}"
        f"\n  Attacks:       {n_attacks}  ×  {n_eps} ε values  +  "
        f"{'PGD T-sweep (' + str(len(PGD_T_GRID)) + ' points)' if not args.no_t_sweep else 'no T-sweep'}"
        f"\n  Total combos:  {len(combos)} ({n * n_attacks * n_eps * n_models} main + {t_extra} T-sweep)"
        f"\n  Output:        {csv_path}"
        f"\n"
        f"\n  Statistical basis: N={n} → SSIM CI ≈ ±{1.96*0.05/n**0.5:.3f} (σ=0.05, 95%)"
        f"\n"
    )

    if args.dry_run:
        for i, (m, a, e, img, t) in enumerate(combos[:10]):
            tag = f" [T-sweep T={t}]" if t else ""
            print(f"  [{i+1}] {m['name']} × {a['name']} × ε={e*255:.0f}/255{tag} × {img.name}")
        if len(combos) > 10:
            print(f"  ... and {len(combos)-10} more")
        return

    # Write CSV header if new file
    done = completed_keys(csv_path) if args.resume else set()
    if done:
        print(f"  Resuming: {len(done)} rows already completed\n")

    if not csv_path.exists() or not args.resume:
        with open(csv_path, "w", newline="") as f:
            extra_cols = ["t_sweep"] if not args.no_t_sweep else []
            writer = csv.DictWriter(f, fieldnames=CSV_COLUMNS + extra_cols)
            writer.writeheader()

    model_cache: dict = {}
    extra_cols = ["t_sweep"] if not args.no_t_sweep else []
    all_cols = CSV_COLUMNS + extra_cols

    pbar = tqdm(combos, desc="B-probe", unit="run")
    for model_cfg, attack_def, eps, img_path, t_val in pbar:
        model_name = model_cfg["name"]
        attack_name = attack_def["name"]
        tag = f" T={t_val}" if t_val else ""
        key = (model_name, attack_name, f"{eps:.5f}", str(args.seed), img_path.name)
        if key in done:
            continue

        pbar.set_postfix_str(f"{model_name[:12]}|{attack_name}|ε={eps*255:.0f}/255{tag}")

        # Load model (cached per name)
        if model_name not in model_cache:
            try:
                model, _, _, vit_cfg = load_model_and_labels(model_cfg["path"], None, device=device)
                model_cache[model_name] = (model, vit_cfg)
            except Exception as e:
                tqdm.write(f"  ERROR loading {model_name}: {e}")
                continue

        model, vit_cfg = model_cache[model_name]

        try:
            res = run_single(model, vit_cfg, attack_def, eps, img_path, device, args.seed)
        except Exception as e:
            tqdm.write(f"  ERROR: {model_name} × {attack_name} × ε={eps*255:.0f} × {img_path.name}: {e}")
            continue

        row = {
            "model": model_name,
            "attack": attack_name,
            "epsilon": f"{eps:.5f}",
            "seed": str(args.seed),
            "image": img_path.name,
            "has_mask": "1" if has_mask_lookup.get(img_path.name, False) else "0",
            "orig_pred": res["orig_pred"],
            "orig_conf": f"{res['orig_conf']:.6f}",
            "adv_pred": res["adv_pred"],
            "adv_conf": f"{res['adv_conf']:.6f}",
            "asr": _fmt(res["asr"]),
            "confidence_drop": _fmt(res["confidence_drop"]),
            "topk_drop": _fmt(res.get("topk_drop", float("nan"))),
            "linf": _fmt(res["linf"]),
            "psnr": _fmt(res["psnr"]),
            "ssim": _fmt(res["ssim"]),
            "lpips": _fmt(res["lpips"]),
            "modified_pixels": _fmt(res["modified_pixels"]),
            "time_seconds": f"{res['elapsed']:.2f}",
        }
        if not args.no_t_sweep:
            row["t_sweep"] = str(t_val) if t_val else ""

        with open(csv_path, "a", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=all_cols)
            writer.writerow(row)

    print(f"\nSweep complete → {csv_path}")

    print("\nGenerating plots and summary...")
    generate_plots(csv_path, out_dir)
    generate_summary(csv_path, out_dir)


if __name__ == "__main__":
    main()