"""Archive old (single-model degenerate) SAGA data antes de submeter SAGA-real.

Move/copia as rows com `attack == "SAGA"` de:
  - results/raw/sweep_main_tcc/<model>/results.csv + attention_maps.parquet
  - results/raw/sweep_ablation_tcc/<model>/results.csv + attention_maps.parquet
  - results/raw/probe_eps_curve_tcc/<model>/results.csv + attention_maps.parquet
  - results/raw/probe_eps_curve_tcc_saga_only/<model>/results.csv + attention_maps.parquet

para:
  results/raw/_archive_saga_single_model/<sweep>/<model>/{results.csv,attention_maps.parquet}

Por DEFAULT é DRY-RUN (não modifica nada — só lista o que seria feito). Pra
executar passa `--commit`. Após executar, os CSV/parquets dos sweeps ficam
SEM rows SAGA — SAGA-real sweeps depois preenchem com dados novos.

NÃO É REVERSÍVEL com `--commit`. Recomendado:
1. Rodar DRY (sem --commit) e confirmar contagens
2. Rodar com `--commit` SÓ após smoke test do SAGA-real ter passado
3. Manter `_archive_saga_single_model/` no disco (pequeno) pra forensics

Usage:
    # Dry-run (default)
    python scripts/archive_old_saga_data.py

    # Commit
    python scripts/archive_old_saga_data.py --commit
"""
from __future__ import annotations

import argparse
import shutil
import sys
from pathlib import Path

import pandas as pd


PROJECT_ROOT = Path(__file__).resolve().parent.parent
RAW_DIR = PROJECT_ROOT / "results" / "raw"
ARCHIVE_DIR = RAW_DIR / "_archive_saga_single_model"

SWEEPS = [
    "sweep_main_tcc",
    "sweep_ablation_tcc",
    "probe_eps_curve_tcc",
    "probe_eps_curve_tcc_saga_only",
]


def split_csv(src_csv: Path, archive_csv: Path, commit: bool) -> tuple[int, int]:
    """Move SAGA rows from src_csv to archive_csv. Returns (n_saga, n_kept)."""
    df = pd.read_csv(src_csv)
    if "attack" not in df.columns:
        return 0, len(df)
    saga_df = df[df["attack"] == "SAGA"].copy()
    kept_df = df[df["attack"] != "SAGA"].copy()
    if len(saga_df) == 0:
        return 0, len(kept_df)
    if commit:
        archive_csv.parent.mkdir(parents=True, exist_ok=True)
        saga_df.to_csv(archive_csv, index=False)
        kept_df.to_csv(src_csv, index=False)
    return len(saga_df), len(kept_df)


def split_parquet(src_pq: Path, archive_pq: Path, commit: bool) -> tuple[int, int]:
    """Move SAGA rows from src_pq to archive_pq. Returns (n_saga, n_kept)."""
    if not src_pq.exists():
        return 0, 0
    df = pd.read_parquet(src_pq)
    if "attack" not in df.columns:
        return 0, len(df)
    saga_df = df[df["attack"] == "SAGA"].copy()
    kept_df = df[df["attack"] != "SAGA"].copy()
    if len(saga_df) == 0:
        return 0, len(kept_df)
    if commit:
        archive_pq.parent.mkdir(parents=True, exist_ok=True)
        saga_df.to_parquet(archive_pq, index=False)
        kept_df.to_parquet(src_pq, index=False)
    return len(saga_df), len(kept_df)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--commit", action="store_true",
                    help="Apply changes (default: dry-run)")
    args = ap.parse_args()

    mode = "COMMIT" if args.commit else "DRY-RUN"
    print(f"=" * 70)
    print(f"Archive old SAGA data [{mode}]")
    print(f"Source: {RAW_DIR}")
    print(f"Archive: {ARCHIVE_DIR}")
    print(f"=" * 70)

    total_saga = 0
    total_kept = 0
    for sweep in SWEEPS:
        sweep_dir = RAW_DIR / sweep
        if not sweep_dir.exists():
            print(f"\n  [SKIP] {sweep}/ — directory not found")
            continue
        print(f"\n  {sweep}/")
        for model_dir in sorted(sweep_dir.glob("vit-*_imagenet-1k_seed42")):
            csv_src = model_dir / "results.csv"
            pq_src = model_dir / "attention_maps.parquet"
            if not csv_src.exists():
                continue
            csv_arch = ARCHIVE_DIR / sweep / model_dir.name / "results.csv"
            pq_arch = ARCHIVE_DIR / sweep / model_dir.name / "attention_maps.parquet"
            n_saga_csv, n_kept_csv = split_csv(csv_src, csv_arch, args.commit)
            n_saga_pq, n_kept_pq = split_parquet(pq_src, pq_arch, args.commit)
            print(f"    {model_dir.name}: csv {n_saga_csv} SAGA / {n_kept_csv} kept | "
                  f"parquet {n_saga_pq} SAGA / {n_kept_pq} kept")
            total_saga += n_saga_csv
            total_kept += n_kept_csv

    print(f"\n{'=' * 70}")
    print(f"Totals: {total_saga:,} SAGA rows {'moved to archive' if args.commit else 'WOULD BE moved'}")
    print(f"        {total_kept:,} non-SAGA rows kept in source")
    print(f"{'=' * 70}")
    if not args.commit:
        print("Dry-run only — pass --commit to apply.")
    return 0


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