"""Lab-server batch arena — run the synthcity model suite (and any other backends) across
datasets and record results to the shared leaderboard.

Designed for a GPU/large-disk host. Run INSIDE an env that has: synthcity, data-designer,
sdmetrics, fairlearn, huggingface_hub (see repo README). Point SYNTHCITY_VENV at the same env
so the subprocess runner finds it, and set LEADERBOARD_REPO + HF_TOKEN to sync results to the
public HF-Dataset leaderboard (otherwise results stay in results/arena_results.csv — copy or
push that file manually).

Usage:
    SYNTHCITY_VENV=$CONDA_PREFIX \
    LEADERBOARD_REPO=nnagesh101/memisislabs-leaderboard HF_TOKEN=hf_xxx \
    python scripts/server_arena.py --datasets adult german compas --rows 1000 --train-cap 5000
"""
from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from pipeline import arena, datasets, leaderboard  # noqa: E402
from pipeline.metadata import build_metadata  # noqa: E402

# synthcity plugins matching the TabDiff + fair-tab-diffusion baselines.
DEFAULT_MODELS = ["ctgan", "tvae", "rtvae", "nflow", "ddpm", "arf", "goggle", "great",
                  "dpgan", "pategan", "adsgan", "privbayes", "decaf"]
DEFAULT_DATASETS = ["openml_45040", "adult", "german", "bank", "compas",
                    "shoppers", "magic", "default", "diabetes_pima", "beijing", "news"]


def available_plugins(requested: list[str]) -> list[str]:
    """Filter requested plugin names to those this synthcity install actually provides."""
    try:
        from synthcity.plugins import Plugins
        have = set(Plugins().list())
    except Exception as exc:  # synthcity not importable in this env
        print(f"[warn] could not list synthcity plugins ({exc}); using requested list as-is")
        return requested
    missing = [m for m in requested if m not in have]
    if missing:
        print(f"[warn] not available in this synthcity install, skipping: {missing}")
    return [m for m in requested if m in have]


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--models", nargs="*", default=DEFAULT_MODELS,
                   help="synthcity plugin names (default: the tabular+privacy+fairness suite)")
    p.add_argument("--datasets", nargs="*", default=DEFAULT_DATASETS)
    p.add_argument("--rows", type=int, default=1000, help="synthetic rows per model")
    p.add_argument("--train-cap", type=int, default=5000,
                   help="cap real rows used (fit + eval reference) for tractable deep-model runs")
    args = p.parse_args()

    models = [("synthcity", m) for m in available_plugins(args.models)]
    if not models:
        sys.exit("no runnable models")
    print(f"models: {[m for _, m in models]}")

    run_stamp = time.strftime("%Y%m%d_%H%M")
    for ds in args.datasets:
        real, target, protected, task = datasets.load(ds)
        if args.train_cap and len(real) > args.train_cap:
            real = real.sample(args.train_cap, random_state=0).reset_index(drop=True)
        md = build_metadata(real)
        print(f"\n=== {ds} ({real.shape}, task={task}) ===", flush=True)
        entries = arena.run_arena(
            models, real, md, target=target, protected=protected, task=task,
            num_records=args.rows,
            on_progress=lambda b, m, i, t: print(f"  [{i}/{t}] {b}·{m} …", flush=True),
        )
        for e in entries:
            if "metrics" in e:
                m = e["metrics"]
                fair = m.get("fairness") or {}
                print(f"  {e['label']:<26} fid={m.get('overall_score'):.3f} "
                      f"priv={m.get('new_row_synthesis')} util={m.get('ml_efficacy')} "
                      f"fair={fair.get('fairness_score')}")
            else:
                print(f"  {e['label']:<26} ERROR: {e.get('error', '?')[:120]}")
        leaderboard.record(entries, dataset=ds, run_id=f"server_{run_stamp}_{ds}",
                           num_records=args.rows)
        print(f"  recorded -> leaderboard ({ds})", flush=True)

    print("\nDONE — results in results/arena_results.csv"
          " (synced to LEADERBOARD_REPO if configured)")


if __name__ == "__main__":
    main()