#!/usr/bin/env python3
"""Merge dos CSVs do probe ε × métricas em um único CSV consolidado.

Concatena os CSVs de 2 sweeps:
  1. `probe_eps_curve_tcc/` — sweep original (4 modelos × 5 ataques × 7 ε × 200 imgs)
     - vit-s-16, vit-s-32: completos com 5 ataques
     - vit-b-32, vit-b-16: parciais (FGSM/PGD/MIM completos + TGR parcial, sem SAGA)
  2. `probe_eps_curve_tcc_saga_only/` — sweep complementar (B/32 e B/16 SAGA)

Filtra TGR (descartado pelo TCC) e gera CSV final com 4 ataques × 4 modelos × 7 ε × 200 imgs.

Adiciona coluna `has_mask` consultando metadata.json do dataset híbrido (se disponível).

Usage:
    python scripts/merge_probe_csvs.py \\
        --original-dir results/raw/probe_eps_curve_tcc \\
        --saga-only-dir results/raw/probe_eps_curve_tcc_saga_only \\
        --metadata data/in1k_hybrid_1k_subset200/metadata.json \\
        --out results/raw/probe_merged_no_tgr.csv
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path


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()


def _load_has_mask_lookup(metadata_path: Path) -> dict[str, bool]:
    """Lê metadata.json do híbrido e retorna {filename: has_mask}."""
    if not metadata_path.exists():
        print(f"  WARN: metadata.json não encontrado em {metadata_path}; "
              f"coluna has_mask ficará vazia")
        return {}
    payload = json.loads(metadata_path.read_text())
    samples = payload.get("samples") or []
    return {s["filename"]: bool(s.get("has_mask", False)) for s in samples}


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--original-dir", type=Path,
                        default=PROJECT_ROOT / "results" / "raw" / "probe_eps_curve_tcc",
                        help="Dir do sweep original (4 subdirs por modelo).")
    parser.add_argument("--saga-only-dir", type=Path,
                        default=PROJECT_ROOT / "results" / "raw" / "probe_eps_curve_tcc_saga_only",
                        help="Dir do sweep saga-only (2 subdirs B/32 e B/16).")
    parser.add_argument("--metadata", type=Path,
                        default=PROJECT_ROOT / "data" / "in1k_hybrid_1k_subset200" / "metadata.json",
                        help="metadata.json do dataset híbrido (pra has_mask).")
    parser.add_argument("--out", type=Path,
                        default=PROJECT_ROOT / "results" / "raw" / "probe_merged_no_tgr.csv",
                        help="CSV de saída.")
    parser.add_argument("--keep-tgr", action="store_true",
                        help="Manter rows TGR (default: filtra).")
    args = parser.parse_args()

    try:
        import pandas as pd
    except ImportError:
        print("ERROR: pandas necessário")
        return 1

    print(f"Merging CSVs:")
    print(f"  original:  {args.original_dir}")
    print(f"  saga_only: {args.saga_only_dir}")

    # Coletar todos os results.csv
    csvs_original = sorted(args.original_dir.glob("*/results.csv"))
    csvs_saga = sorted(args.saga_only_dir.glob("*/results.csv")) \
        if args.saga_only_dir.exists() else []

    print(f"\nCSVs encontrados:")
    print(f"  original:  {len(csvs_original)}")
    for c in csvs_original:
        n = sum(1 for _ in open(c)) - 1
        print(f"    {c.parent.name}: {n} rows")
    print(f"  saga_only: {len(csvs_saga)}")
    for c in csvs_saga:
        n = sum(1 for _ in open(c)) - 1
        print(f"    {c.parent.name}: {n} rows")

    if not csvs_original:
        print(f"\nERROR: nenhum CSV em {args.original_dir}")
        return 1

    # Concat
    dfs = [pd.read_csv(c) for c in csvs_original + csvs_saga]
    df = pd.concat(dfs, ignore_index=True)
    print(f"\nTotal pré-filtro: {len(df)} rows")
    print(f"  ataques: {sorted(df['attack'].unique())}")
    print(f"  modelos: {sorted(df['model'].unique())}")
    print(f"  ε: {sorted(df['epsilon'].unique())}")

    # Filtrar TGR (a menos que --keep-tgr)
    if not args.keep_tgr:
        n_tgr = (df['attack'] == 'TGR').sum()
        df = df[df['attack'] != 'TGR'].reset_index(drop=True)
        print(f"\nFiltrado TGR: {n_tgr} rows removidas → {len(df)} rows")

    # Adicionar has_mask via metadata
    has_mask_lookup = _load_has_mask_lookup(args.metadata)
    if has_mask_lookup:
        df['has_mask'] = df['image'].map(has_mask_lookup).fillna(False).astype(int)
        n_mask = df['has_mask'].sum()
        print(f"\nhas_mask: {n_mask} rows com mask, {len(df) - n_mask} sem")

    # Adicionar coluna eps_255 conveniente
    df['eps_255'] = (df['epsilon'].astype(float) * 255).round().astype(int)

    # Sanity: dedup por (model, attack, epsilon, image, seed) — caso CSVs se sobreponham
    key_cols = ['model', 'attack', 'epsilon', 'image', 'seed']
    n_before = len(df)
    df = df.drop_duplicates(subset=key_cols, keep='first').reset_index(drop=True)
    n_dups = n_before - len(df)
    if n_dups > 0:
        print(f"\nDedup: {n_dups} rows duplicadas removidas → {len(df)} rows")

    # Reportar contagem final por (model, attack)
    print(f"\nContagem por (modelo × ataque):")
    counts = df.groupby(['model', 'attack']).size().unstack(fill_value=0)
    print(counts.to_string())

    # Save
    args.out.parent.mkdir(parents=True, exist_ok=True)
    df.to_csv(args.out, index=False)
    print(f"\n✓ Salvo: {args.out}")
    print(f"  Total: {len(df)} rows | "
          f"{df['model'].nunique()} modelos × "
          f"{df['attack'].nunique()} ataques × "
          f"{df['eps_255'].nunique()} ε × "
          f"{df['image'].nunique()} imgs")

    # Esperado: 4 × 4 × 7 × 200 = 22.400 rows
    expected = (df['model'].nunique() * df['attack'].nunique()
                * df['eps_255'].nunique() * df['image'].nunique())
    if len(df) == expected:
        print(f"  ✓ Bate com esperado ({expected})")
    else:
        print(f"  ⚠️ Esperado {expected} rows mas tem {len(df)} — alguma combinação faltando?")

    return 0


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