"""Verify every numerical claim in paper/main.tex against raw sweep data.

Reads:
- results/raw/sweep_main_tcc/<arch>/results.csv (4 architectures)
- results/raw/sweep_main_tcc/_emd_augmented.parquet (W_1 = attention_emd_2d)
- results/raw/sweep_main_tcc/fg_trajectory_modes.parquet (fg_mass trajectories)

Reports every assertion that fails or is outside a documented tolerance.

Usage:
    cd /Users/ldmc/Desktop/faculdade/serrapilheira/ViTViz
    python scripts/verify_paper_numbers.py
"""
from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import pandas as pd
from scipy.stats import spearmanr

PROJECT_ROOT = Path(__file__).resolve().parent.parent
RAW = PROJECT_ROOT / "results" / "raw" / "sweep_main_tcc"

# Tolerance for floating-point comparisons of correlations and ratios
TOL_RHO = 0.005
TOL_MEAN = 0.0010
TOL_INT = 0  # exact for counts

results: list[tuple[str, str, str]] = []  # (status, label, message)


def check(label: str, computed: float, claimed: float, tol: float = TOL_RHO) -> None:
    diff = abs(computed - claimed)
    if diff <= tol:
        results.append(("OK", label, f"{computed:.4f} vs {claimed:.4f} (Δ={diff:.4f}, tol={tol})"))
    else:
        results.append(("FAIL", label, f"{computed:.4f} vs claimed {claimed:.4f} (Δ={diff:.4f}, tol={tol})"))


def check_int(label: str, computed: int, claimed: int) -> None:
    if computed == claimed:
        results.append(("OK", label, f"{computed} (exact)"))
    else:
        results.append(("FAIL", label, f"{computed} vs claimed {claimed}"))


def load_data() -> pd.DataFrame:
    archs = {
        "ViT-S/16": "vit-s-16",
        "ViT-S/32": "vit-s-32",
        "ViT-B/32": "vit-b-32",
        "ViT-B/16": "vit-b-16",
    }
    csv_all = []
    for label, sub in archs.items():
        csv = pd.read_csv(RAW / f"{sub}_imagenet-1k_seed42" / "results.csv")
        csv["model"] = label
        csv_all.append(csv)
    csv_df = pd.concat(csv_all, ignore_index=True)
    emd = pd.read_parquet(RAW / "_emd_augmented.parquet")
    df = csv_df.merge(
        emd[["model", "attack", "image", "attention_emd_2d"]],
        on=["model", "attack", "image"],
        how="left",
    )
    df["abs_dh"] = df["attention_entropy_delta"].abs()
    df["W1"] = df["attention_emd_2d"]
    df = df.dropna(subset=["W1"]).reset_index(drop=True)
    return df


def verify_section_v_b(df: pd.DataFrame) -> None:
    """§V.B Cross-Attack: pooled per-arch W1, JSD, |dH| on iteratives."""
    print("\n=== §V.B Cross-Attack: pooled per-arch (iteratives) ===")
    iters = ["PGD", "MIM", "SAGA"]
    df_i = df[df.attack.isin(iters)]

    # Existing paper claims at line 798-799 (pre-edit values)
    claims_w1 = {
        "ViT-S/16": 0.0724,
        "ViT-S/32": 0.0909,
        "ViT-B/32": 0.0802,
        "ViT-B/16": 0.0860,
    }
    claims_jsd = {
        "ViT-S/16": 0.226,
        "ViT-S/32": 0.202,
        "ViT-B/32": 0.194,
        "ViT-B/16": 0.195,
    }
    claims_dh = {
        "ViT-S/16": 0.225,
        "ViT-S/32": 0.248,
        "ViT-B/32": 0.197,
        "ViT-B/16": 0.303,
    }
    for arch in ["ViT-S/16", "ViT-S/32", "ViT-B/32", "ViT-B/16"]:
        sub = df_i[df_i.model == arch]
        check(f"§V.B W1 pool {arch}", sub.W1.mean(), claims_w1[arch], tol=TOL_MEAN)
        check(f"§V.B JSD pool {arch}", sub.attention_jsd.mean(), claims_jsd[arch], tol=TOL_MEAN)
        check(f"§V.B |dH| pool {arch}", sub.abs_dh.mean(), claims_dh[arch], tol=TOL_MEAN)


def verify_section_iv_d_3(df: pd.DataFrame) -> None:
    """§IV.D.3: cross-axis correlations and 5x5 cross-block."""
    print("\n=== §IV.D.3 Cross-axis correlations ===")
    df_clean = df.dropna(subset=["W1", "attention_jsd", "abs_dh"])

    # Existing paper claim: full-sweep correlations (line 658)
    r, _ = spearmanr(df_clean.W1, df_clean.attention_jsd)
    check("§IV.D.3 ρ(W1, JSD) full sweep", r, 0.62, tol=TOL_RHO + 0.005)
    r, _ = spearmanr(df_clean.W1, df_clean.abs_dh)
    check("§IV.D.3 ρ(W1, |dH|) full sweep", r, 0.31, tol=TOL_RHO + 0.005)
    r, _ = spearmanr(df_clean.attention_jsd, df_clean.abs_dh)
    check("§IV.D.3 ρ(JSD, |dH|) full sweep", r, 0.19, tol=TOL_RHO + 0.005)

    # Off-diagonal stats (from line ~668)
    q1_w1, q3_w1 = df_clean.W1.quantile([0.25, 0.75]).values
    q1_jsd, q3_jsd = df_clean.attention_jsd.quantile([0.25, 0.75]).values
    low_high = df_clean[(df_clean.W1 <= q1_w1) & (df_clean.attention_jsd >= q3_jsd)]
    high_low = df_clean[(df_clean.W1 >= q3_w1) & (df_clean.attention_jsd <= q1_jsd)]
    total_off = len(low_high) + len(high_low)
    check_int("§IV.D.3 n(low-W1 × high-JSD)", len(low_high), 140)
    check_int("§IV.D.3 n(high-W1 × low-JSD)", len(high_low), 116)
    check_int("§IV.D.3 total off-diagonal", total_off, 256)
    check("§IV.D.3 off-diagonal fraction", total_off / len(df_clean), 0.0160, tol=0.0005)

    # 5×5 cross-block on iterative + Guillaumin
    print("\n=== §IV.D.3 5×5 cross-block (iterative + Guillaumin) ===")
    iters = ["PGD", "MIM", "SAGA"]
    df_c = df[df.attack.isin(iters)].dropna(
        subset=["miou_adv", "map_adv", "W1", "attention_jsd", "abs_dh"]
    )
    check_int("§IV.D.3 cross-block n", len(df_c), 1140)
    pairs = [
        ("W1", "miou_adv", "W1 × mIoU", -0.056),
        ("W1", "map_adv", "W1 × mAP", -0.026),
        ("attention_jsd", "miou_adv", "JSD × mIoU", -0.136),
        ("attention_jsd", "map_adv", "JSD × mAP", -0.078),
        ("abs_dh", "miou_adv", "|dH| × mIoU", +0.022),
        ("abs_dh", "map_adv", "|dH| × mAP", +0.045),
        ("miou_adv", "map_adv", "mIoU × mAP", +0.536),
    ]
    for c1, c2, label, claimed in pairs:
        r, _ = spearmanr(df_c[c1], df_c[c2])
        check(f"§IV.D.3 ρ({label})", r, claimed)


def verify_section_v_c(df: pd.DataFrame) -> None:
    """§V.C ρ(conf_drop, axis) per attack."""
    print("\n=== §V.C ρ(conf_drop, axis) per attack ===")
    claims = {
        # attack: (W1, JSD, |dH|)
        "FGSM": (0.050, 0.083, 0.044),
        "PGD": (0.076, 0.131, 0.089),
        "MIM": (0.087, 0.130, 0.072),
        "SAGA": (0.169, 0.193, 0.128),
    }
    for atk, (cw, cj, cd) in claims.items():
        sub = df[df.attack == atk].dropna(
            subset=["confidence_drop", "W1", "attention_jsd", "abs_dh"]
        )
        rw, _ = spearmanr(sub.confidence_drop, sub.W1)
        rj, _ = spearmanr(sub.confidence_drop, sub.attention_jsd)
        rd, _ = spearmanr(sub.confidence_drop, sub.abs_dh)
        check(f"§V.C {atk} ρ(W1)", rw, cw)
        check(f"§V.C {atk} ρ(JSD)", rj, cj)
        check(f"§V.C {atk} ρ(|dH|)", rd, cd)


def verify_section_v_e(df: pd.DataFrame) -> None:
    """§V.E ΔmIoU per cell + Δfg_mass S-row/B-row split."""
    print("\n=== §V.E Endpoint foreground drainage ===")
    df_g = df.dropna(subset=["miou_clean", "miou_adv"]).copy()
    df_g["d_miou"] = df_g.miou_adv - df_g.miou_clean
    cell = df_g.groupby(["model", "attack"])["d_miou"].mean().reset_index()
    n_neg = (cell.d_miou < 0).sum()
    check_int("§V.E n(cells ΔmIoU < 0)", n_neg, 12)
    check_int("§V.E total cells", len(cell), 16)

    # Three most negative cells
    sorted_cell = cell.sort_values("d_miou")
    top3 = sorted_cell.head(3)
    # Paper claim: S/32×SAGA, S/16×MIM, S/16×PGD
    expected_top3 = [
        ("ViT-S/32", "SAGA", -0.026),
        ("ViT-S/16", "MIM", -0.025),
        ("ViT-S/16", "PGD", -0.025),
    ]
    for i, (m, a, v) in enumerate(expected_top3):
        row = top3.iloc[i]
        ok_m = row.model == m
        ok_a = row.attack == a
        if ok_m and ok_a:
            check(f"§V.E ΔmIoU #{i+1} {m}×{a}", row.d_miou, v, tol=0.0015)
        else:
            results.append(
                (
                    "FAIL",
                    f"§V.E ΔmIoU rank #{i+1}",
                    f"expected {m}×{a}, got {row.model}×{row.attack} ({row.d_miou:+.4f})",
                )
            )

    # Specific outlier cells claimed
    bfgsm_b32 = cell[(cell.model == "ViT-B/32") & (cell.attack == "FGSM")].d_miou.values[0]
    check("§V.E B/32×FGSM ΔmIoU", bfgsm_b32, +0.028, tol=0.0015)
    s16_fgsm = cell[(cell.model == "ViT-S/16") & (cell.attack == "FGSM")].d_miou.values[0]
    check("§V.E S/16×FGSM ΔmIoU", s16_fgsm, +0.014, tol=0.0015)
    b16_fgsm = cell[(cell.model == "ViT-B/16") & (cell.attack == "FGSM")].d_miou.values[0]
    check("§V.E B/16×FGSM ΔmIoU", b16_fgsm, -0.018, tol=0.0015)

    # S-row vs B-row iterative pool
    iters = ["PGD", "MIM", "SAGA"]
    s_pool = df_g[df_g.model.isin(["ViT-S/16", "ViT-S/32"]) & df_g.attack.isin(iters)].d_miou.mean()
    b_pool = df_g[df_g.model.isin(["ViT-B/16", "ViT-B/32"]) & df_g.attack.isin(iters)].d_miou.mean()
    check("§V.E S-row iterative ΔmIoU pool", s_pool, -0.0225, tol=0.0015)
    check("§V.E B-row iterative ΔmIoU pool", b_pool, -0.0094, tol=0.0015)

    # Binomial test for 12/16 against p=0.5
    from scipy.stats import binomtest
    p_value = binomtest(12, 16, p=0.5, alternative="two-sided").pvalue
    check("§V.E binomial p(12/16 vs 0.5)", p_value, 0.077, tol=0.005)

    # Per-cell bootstrap CIs (5 strictly negative, 1 strictly positive, 10 span zero)
    rng = np.random.default_rng(42)
    n_neg_strict, n_pos_strict, n_span = 0, 0, 0
    for (_m, _a), sub in df_g.groupby(["model", "attack"]):
        arr = sub.d_miou.values
        boots = [arr[rng.integers(0, len(arr), len(arr))].mean() for _ in range(1000)]
        lo, hi = np.percentile(boots, [2.5, 97.5])
        if hi < 0:
            n_neg_strict += 1
        elif lo > 0:
            n_pos_strict += 1
        else:
            n_span += 1
    check_int("§V.E cells CI strictly negative", n_neg_strict, 5)
    check_int("§V.E cells CI strictly positive", n_pos_strict, 1)
    check_int("§V.E cells CI span zero", n_span, 10)

    # Δfg_mass per cell from parquet
    print("\n=== §V.E Δfg_mass per cell ===")
    fg = pd.read_parquet(RAW / "fg_trajectory_modes.parquet")
    fg["arch_label"] = fg.model.str.replace("ViT-", "", regex=False)
    fg_cell = fg.groupby(["arch_label", "attack"])["delta"].mean().reset_index()
    # S-row positive, B-row mixed (with B/32 negative)
    s16_dfg = fg_cell[(fg_cell.arch_label == "S/16") & fg_cell.attack.isin(iters)].delta.values
    s32_dfg = fg_cell[(fg_cell.arch_label == "S/32") & fg_cell.attack.isin(iters)].delta.values
    b16_dfg = fg_cell[(fg_cell.arch_label == "B/16") & fg_cell.attack.isin(iters)].delta.values
    b32_dfg = fg_cell[(fg_cell.arch_label == "B/32") & fg_cell.attack.isin(iters)].delta.values
    # All S-row Δfg_mass > 0
    if all(v > 0 for v in s16_dfg) and all(v > 0 for v in s32_dfg):
        results.append(("OK", "§V.E S-row Δfg_mass all positive", f"S/16={s16_dfg}, S/32={s32_dfg}"))
    else:
        results.append(("FAIL", "§V.E S-row Δfg_mass all positive", f"S/16={s16_dfg}, S/32={s32_dfg}"))
    # All B/32 Δfg_mass < 0
    if all(v < 0 for v in b32_dfg):
        results.append(("OK", "§V.E B/32 Δfg_mass all negative", f"{b32_dfg}"))
    else:
        results.append(("FAIL", "§V.E B/32 Δfg_mass all negative", f"{b32_dfg}"))


def verify_section_vi_b(df: pd.DataFrame) -> None:
    """§VI.B pixel × attention max ρ and threshold counts."""
    print("\n=== §VI.B Pixel × attention correlations ===")
    all_rhos = []
    for atk in ["FGSM", "PGD", "MIM", "SAGA"]:
        sub = df[df.attack == atk].dropna(
            subset=["psnr", "ssim", "lpips", "W1", "attention_jsd", "abs_dh"]
        )
        for px in ["psnr", "ssim", "lpips"]:
            for ax in ["W1", "attention_jsd", "abs_dh"]:
                r, _ = spearmanr(sub[px], sub[ax])
                all_rhos.append((atk, px, ax, r))

    max_abs = max(abs(r) for _, _, _, r in all_rhos)
    check("§VI.B max |ρ| pixel × attention", max_abs, 0.216, tol=0.005)

    # Specific max cell
    abs_sorted = sorted(all_rhos, key=lambda x: -abs(x[3]))
    top = abs_sorted[0]
    if top[0] == "SAGA" and top[1] == "lpips" and top[2] == "attention_jsd":
        results.append(("OK", "§VI.B max cell (SAGA × LPIPS × JSD)", f"ρ={top[3]:+.4f}"))
    else:
        results.append(
            ("FAIL", "§VI.B max cell", f"expected SAGA×LPIPS×JSD, got {top[0]}×{top[1]}×{top[2]}")
        )

    n_above_015 = sum(1 for _, _, _, r in all_rhos if abs(r) > 0.15)
    n_above_020 = sum(1 for _, _, _, r in all_rhos if abs(r) > 0.20)
    check_int("§VI.B n cells |ρ|>0.15", n_above_015, 6)
    check_int("§VI.B n cells |ρ|>0.20", n_above_020, 2)

    # Holm-Bonferroni count of significant cells
    pvals = []
    for atk in ["FGSM", "PGD", "MIM", "SAGA"]:
        sub = df[df.attack == atk].dropna(
            subset=["psnr", "ssim", "lpips", "W1", "attention_jsd", "abs_dh"]
        )
        for px in ["psnr", "ssim", "lpips"]:
            for ax in ["W1", "attention_jsd", "abs_dh"]:
                _, p = spearmanr(sub[px], sub[ax])
                pvals.append(p)
    pvals_sorted = sorted(pvals)
    m = len(pvals_sorted)
    n_sig = sum(1 for i, p in enumerate(pvals_sorted) if p < 0.05 / (m - i))
    check_int("§VI.B Holm-Bonferroni significant cells", n_sig, 21)


def verify_asr_and_success_failure(df: pd.DataFrame) -> None:
    """ASR per cell + success/failure W1 ratios for the key claims in §V.B and §V.C."""
    print("\n=== ASR per cell + SAGA/FGSM success-failure ratios ===")
    # Per-cell ASR (paper claims spread across §V.B abstract and Table)
    asr_claims = {
        ("ViT-B/32", "PGD"): 0.920,
        ("ViT-B/32", "SAGA"): 0.782,
        ("ViT-B/16", "FGSM"): 0.674,
        ("ViT-S/16", "PGD"): 0.941,
        ("ViT-S/16", "SAGA"): 0.883,
    }
    for (m, a), claimed in asr_claims.items():
        sub = df[(df.model == m) & (df.attack == a)]
        check(f"§V.B ASR {m} × {a}", sub.asr.mean(), claimed, tol=0.005)

    # Pooled ASR by attack
    pooled_asr_claims = {"PGD": 0.931, "MIM": 0.926, "SAGA": 0.873}
    for atk, claimed in pooled_asr_claims.items():
        check(f"§V.B pooled ASR {atk}", df[df.attack == atk].asr.mean(), claimed, tol=0.005)

    # PGD-SAGA gap (paper claims 5.8 pp pooled, 13.8 pp on B/32)
    gap_pooled = (df[df.attack == "PGD"].asr.mean() - df[df.attack == "SAGA"].asr.mean()) * 100
    check("§V.B PGD-SAGA gap (pp)", gap_pooled, 5.8, tol=0.2)

    gap_b32 = (
        df[(df.model == "ViT-B/32") & (df.attack == "PGD")].asr.mean()
        - df[(df.model == "ViT-B/32") & (df.attack == "SAGA")].asr.mean()
    ) * 100
    check("§V.B PGD-SAGA gap B/32 (pp)", gap_b32, 13.8, tol=0.2)

    # SAGA success/failure W1 ratio pooled = 1.33×, B/32 = 1.64× (paper §V.C)
    saga = df[df.attack == "SAGA"].dropna(subset=["W1"])
    ratio_saga = saga[saga.asr == 1].W1.mean() / saga[saga.asr == 0].W1.mean()
    check("§V.C SAGA succ/fail W1 ratio", ratio_saga, 1.33, tol=0.02)

    saga_b32 = df[(df.model == "ViT-B/32") & (df.attack == "SAGA")].dropna(subset=["W1"])
    ratio_b32 = saga_b32[saga_b32.asr == 1].W1.mean() / saga_b32[saga_b32.asr == 0].W1.mean()
    check("§V.C SAGA succ/fail W1 ratio B/32", ratio_b32, 1.643, tol=0.02)

    # FGSM inverse separation: succ < fail (paper says 0.91×)
    fgsm = df[df.attack == "FGSM"].dropna(subset=["W1"])
    ratio_fgsm = fgsm[fgsm.asr == 1].W1.mean() / fgsm[fgsm.asr == 0].W1.mean()
    check("§V.C FGSM succ/fail W1 ratio", ratio_fgsm, 0.91, tol=0.02)


def verify_baseline(df: pd.DataFrame) -> None:
    """§V.A baseline: clean accuracies, mIoU, H_norm pair structure."""
    print("\n=== §V.A Baseline (clean acc, mIoU, H_norm) ===")
    # clean accuracies from FGSM rows (any attack works for orig_pred, but cleanest is FGSM/PGD)
    # The paper says: 0.756 S/16, 0.678 S/32, 0.764 B/32, 0.800 B/16
    df_clean = df[df.attack == "FGSM"].copy()
    df_clean["correct"] = df_clean.orig_pred == df_clean.ground_truth
    claims_acc = {"ViT-S/16": 0.756, "ViT-S/32": 0.678, "ViT-B/32": 0.764, "ViT-B/16": 0.800}
    for arch, claimed in claims_acc.items():
        sub = df_clean[df_clean.model == arch]
        acc = sub.correct.mean()
        check(f"§V.A clean acc {arch}", acc, claimed, tol=0.005)

    # Clean mIoU (only rows with masks)
    df_g = df_clean.dropna(subset=["miou_clean"])
    claims_miou = {"ViT-S/16": 0.177, "ViT-S/32": 0.291, "ViT-B/32": 0.229, "ViT-B/16": 0.285}
    for arch, claimed in claims_miou.items():
        sub = df_g[df_g.model == arch]
        check(f"§V.A clean mIoU {arch}", sub.miou_clean.mean(), claimed, tol=0.005)


def main() -> int:
    print("Loading data...")
    df = load_data()
    print(f"  Loaded {len(df)} rows ({len(df.model.unique())} archs × {len(df.attack.unique())} attacks)")

    verify_baseline(df)
    verify_asr_and_success_failure(df)
    verify_section_iv_d_3(df)
    verify_section_v_b(df)
    verify_section_v_c(df)
    verify_section_v_e(df)
    verify_section_vi_b(df)

    print("\n" + "=" * 70)
    print("VERIFICATION REPORT")
    print("=" * 70)
    n_ok = sum(1 for s, _, _ in results if s == "OK")
    n_fail = sum(1 for s, _, _ in results if s == "FAIL")
    for status, label, msg in results:
        marker = "✓" if status == "OK" else "✗"
        print(f"[{marker}] {label:<50} {msg}")
    print("=" * 70)
    print(f"TOTAL: {n_ok} OK, {n_fail} FAIL  (of {len(results)} checks)")
    return 0 if n_fail == 0 else 1


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