"""Benchmark FelaTab vs TabFM vs XGBoost/LightGBM on OpenML classification datasets.

FelaTab + TabFM run zero-shot / in-context (train split = support rows, capped).
XGBoost / LightGBM are trained on the full train split (CPU; pip builds can't use ROCm).

Outputs:
  benchmark/results.csv            - per (dataset, model) metrics + efficiency
  benchmark/model_card_snippet.md  - Hugging Face model-card-ready Markdown tables

Usage:
  python benchmark.py --smoke                 # 2 tiny datasets, 3 models, minutes
  python benchmark.py                         # full battery
  python benchmark.py --device cpu --skip tabfm
"""

from __future__ import annotations

import argparse
import os
import sys
import time
import tracemalloc
import warnings
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np
import pandas as pd
import psutil

warnings.filterwarnings("ignore")

SEED = 42
HERE = Path(__file__).resolve().parent
# works both as <workspace>/benchmark/benchmark.py and inside the fela-tab repo as
# fela-tab/benchmark/benchmark.py
for _cand in (HERE.parent / "fela-tab", HERE.parent, HERE / ".." / "fela-tab"):
    if (_cand / "modeling.py").is_file():
        FELA_DIR = _cand.resolve()
        break
else:
    FELA_DIR = HERE.parent / "fela-tab"

# --------------------------------------------------------------------------
# Datasets
# --------------------------------------------------------------------------
# (display name, openml name, openml version)
FULL_DATASETS = [
    ("adult", "adult", 2),
    ("credit-g", "credit-g", 1),
    ("blood-transfusion", "blood-transfusion-service-center", 1),
    ("churn", "churn", 1),
    ("electricity", "electricity", 1),
    ("vehicle", "vehicle", 1),
    ("segment", "segment", 1),
    ("jungle_chess", "jungle_chess_2pcs_raw_endgame_complete", 1),
]
SMOKE_DATASETS = [
    ("breast-cancer", "wdbc", 1),
    ("vehicle", "vehicle", 1),
]
MAX_CLASSES = 10  # FelaTab / TabFM constraint


def load_openml(name: str, version: int):
    """Return (X DataFrame, y int array, n_classes). Encodes categoricals, imputes."""
    from sklearn.datasets import fetch_openml

    ds = fetch_openml(name, version=version, as_frame=True, parser="auto")
    X, y = ds.data, ds.target
    # encode target
    y = pd.Categorical(y)
    y_int = y.codes.astype(np.int64)
    n_classes = len(y.categories)
    # encode features
    X = X.copy()
    for c in X.columns:
        if not pd.api.types.is_numeric_dtype(X[c]):
            X[c] = pd.Categorical(X[c]).codes.astype(np.float64)
    X = X.astype(np.float64)
    # impute with column median
    X = X.fillna(X.median(numeric_only=True)).fillna(0.0)
    return X, y_int, n_classes


# --------------------------------------------------------------------------
# Efficiency measurement helpers
# --------------------------------------------------------------------------
class Meter:
    """Tracks wall time, peak process RAM (tracemalloc + psutil RSS), peak VRAM."""

    def __init__(self, device: str):
        self.device = device
        self.proc = psutil.Process(os.getpid())
        self._torch = None
        if device.startswith("cuda"):
            try:
                import torch

                self._torch = torch
                torch.cuda.reset_peak_memory_stats()
            except Exception:
                pass

    def __enter__(self):
        tracemalloc.start()
        self._rss0 = self.proc.memory_info().rss
        self._peak_rss = self._rss0
        self.t0 = time.perf_counter()
        return self

    def sample(self):
        self._peak_rss = max(self._peak_rss, self.proc.memory_info().rss)

    def __exit__(self, *exc):
        self.elapsed = time.perf_counter() - self.t0
        _, py_peak = tracemalloc.get_traced_memory()
        tracemalloc.stop()
        self.sample()
        self.peak_ram_mb = max(py_peak, self._peak_rss - self._rss0) / 1e6
        self.peak_vram_mb = (
            self._torch.cuda.max_memory_allocated() / 1e6 if self._torch else None
        )
        return False


@dataclass
class Result:
    dataset: str
    model: str
    device: str
    roc_auc: float = np.nan
    log_loss: float = np.nan
    accuracy: float = np.nan
    f1_macro: float = np.nan
    fit_s: float = np.nan
    latency_ms: float = np.nan  # per sample
    peak_ram_mb: float = np.nan
    peak_vram_mb: float = np.nan
    status: str = "ok"
    notes: str = ""


# --------------------------------------------------------------------------
# Model adapters: uniform fit_predict_proba(Xtr, ytr, Xte) -> (proba, fit_s, infer_s)
# proba shape [n_test, n_classes], columns aligned with class index 0..K-1
# --------------------------------------------------------------------------
class Adapter:
    name = "base"
    device = "cpu"

    def fit_predict_proba(self, Xtr, ytr, Xte, n_classes):
        raise NotImplementedError


class FelaTabAdapter(Adapter):
    def __init__(self, tier: str, device: str):
        import torch

        sys.path.insert(0, str(FELA_DIR))
        from modeling import load_model  # fela-tab/modeling.py

        self.tier = tier
        self.name = f"FelaTab-{tier}"
        self.device = "gpu" if device.startswith("cuda") else "cpu"
        dev = torch.device("cuda" if device.startswith("cuda") else "cpu")
        self._predict_mod = __import__("modeling")
        self.model = load_model(str(FELA_DIR), tier=tier).to(dev)

    def fit_predict_proba(self, Xtr, ytr, Xte, n_classes):
        from modeling import predict

        t0 = time.perf_counter()
        proba = predict(
            self.model, Xtr, ytr, Xte,
            task="classification", n_classes=n_classes, support_cap=3000,
        )
        total = time.perf_counter() - t0
        # zero-shot: no fit; count full in-context forward as inference,
        # report support ingestion separately via fit_s=0
        return proba, 0.0, total


class TabFMAdapter(Adapter):
    """google/tabfm-1.0.0-pytorch zero-shot via the `tabfm` package.

    Requires HF_TOKEN with accepted license (non-commercial, tabfm-non-commercial-v1.0).
    Sklearn-style: TabFMClassifier.fit(X_train, y_train) -> predict_proba(X_test).
    """

    name = "TabFM"

    def __init__(self, device: str):
        from huggingface_hub import model_info

        self.device = "gpu" if device.startswith("cuda") else "cpu"
        token = os.environ.get("HF_TOKEN")
        if not token:
            raise RuntimeError(
                "HF_TOKEN not set; google/tabfm-1.0.0-pytorch is a gated repo")
        model_info("google/tabfm-1.0.0-pytorch", token=token)  # raises if unauthorized
        import torch
        from tabfm import TabFMClassifier, tabfm_v1_0_0_pytorch as tabfm_ckpt

        self._torch = torch
        self._dev = torch.device("cuda" if device.startswith("cuda") else "cpu")
        model = tabfm_ckpt.load(model_type="classification")
        try:
            model = model.to(self._dev)
        except Exception:
            pass
        self.clf = TabFMClassifier(model=model)

    def fit_predict_proba(self, Xtr, ytr, Xte, n_classes):
        if n_classes > MAX_CLASSES:
            raise RuntimeError(f"TabFM supports <= {MAX_CLASSES} classes")
        t0 = time.perf_counter()
        self.clf.fit(Xtr, ytr)  # in-context: stores support rows, no training
        fit_s = time.perf_counter() - t0
        t1 = time.perf_counter()
        proba = np.asarray(self.clf.predict_proba(Xte), dtype=np.float64)
        infer_s = time.perf_counter() - t1
        return proba[:, :n_classes], fit_s, infer_s


class SklearnAdapter(Adapter):
    def __init__(self, model, name):
        self.model = model
        self.name = name
        self.device = "cpu"

    def fit_predict_proba(self, Xtr, ytr, Xte, n_classes):
        t0 = time.perf_counter()
        self.model.fit(Xtr, ytr)
        fit_s = time.perf_counter() - t0
        t1 = time.perf_counter()
        proba = self.model.predict_proba(Xte)
        infer_s = time.perf_counter() - t1
        return proba, fit_s, infer_s


def build_models(device: str, tiers, skip: set) -> list[Adapter]:
    models: list[Adapter] = []
    for tier in tiers:
        try:
            models.append(FelaTabAdapter(tier, device))
            print(f"[load] FelaTab-{tier} on {device}")
        except Exception as e:
            print(f"[skip] FelaTab-{tier}: {e}")
    if "tabfm" not in skip:
        try:
            models.append(TabFMAdapter(device))
            print(f"[load] TabFM on {device}")
        except Exception as e:
            print(f"[skip] TabFM: {e}")
    if "xgb" not in skip:
        from xgboost import XGBClassifier

        models.append(SklearnAdapter(
            XGBClassifier(n_estimators=300, max_depth=6, learning_rate=0.1,
                          tree_method="hist", n_jobs=-1, random_state=SEED),
            "XGBoost"))
    if "lgbm" not in skip:
        from lightgbm import LGBMClassifier

        models.append(SklearnAdapter(
            LGBMClassifier(n_estimators=300, learning_rate=0.1, n_jobs=-1,
                           random_state=SEED, verbose=-1),
            "LightGBM"))
    return models


# --------------------------------------------------------------------------
# Evaluation
# --------------------------------------------------------------------------
def evaluate(y_true, proba, n_classes):
    from sklearn.metrics import accuracy_score, f1_score, log_loss, roc_auc_score

    proba = np.asarray(proba, dtype=np.float64)
    proba = np.clip(proba, 1e-12, 1.0)
    proba = proba / proba.sum(1, keepdims=True)
    pred = proba.argmax(1)
    acc = accuracy_score(y_true, pred)
    f1 = f1_score(y_true, pred, average="macro")
    labels = list(range(n_classes))
    ll = log_loss(y_true, proba, labels=labels)
    if n_classes == 2:
        auc = roc_auc_score(y_true, proba[:, 1])
    else:
        auc = roc_auc_score(y_true, proba, multi_class="ovr", average="weighted",
                            labels=labels)
    return auc, ll, acc, f1


def run(args) -> list[Result]:
    from sklearn.model_selection import train_test_split

    device = "cuda" if args.device == "gpu" else "cpu"
    if device == "cuda":
        import torch

        if not torch.cuda.is_available():
            print("[warn] GPU requested but unavailable; falling back to CPU")
            device = "cpu"

    dsets = SMOKE_DATASETS if args.smoke else FULL_DATASETS
    if args.datasets:
        keep = set(args.datasets.split(","))
        dsets = [d for d in dsets if d[0] in keep]
    tiers = ["small"] if args.smoke else args.tiers.split(",")
    skip = set(args.skip.split(",")) if args.skip else set()
    if args.smoke:
        skip.add("tabfm")
    models = build_models(device, tiers, skip)
    if not models:
        sys.exit("no models available")

    results: list[Result] = []
    for dname, oml_name, ver in dsets:
        print(f"\n=== {dname} (openml:{oml_name} v{ver}) ===")
        try:
            X, y, n_classes = load_openml(oml_name, ver)
        except Exception as e:
            print(f"  [skip dataset] {e}")
            continue
        if n_classes > MAX_CLASSES:
            print(f"  [skip dataset] {n_classes} classes > {MAX_CLASSES}")
            continue
        Xtr, Xte, ytr, yte = train_test_split(
            X, y, test_size=0.2, random_state=SEED, stratify=y)
        print(f"  train={len(Xtr)} test={len(Xte)} feats={Xtr.shape[1]} classes={n_classes}")
        for m in models:
            r = Result(dataset=dname, model=m.name, device=m.device)
            try:
                with Meter(device if m.device == "gpu" else "cpu") as meter:
                    proba, fit_s, infer_s = m.fit_predict_proba(
                        Xtr.to_numpy(), ytr, Xte.to_numpy(), n_classes)
                r.fit_s = fit_s
                r.latency_ms = 1e3 * infer_s / len(Xte)
                r.peak_ram_mb = meter.peak_ram_mb
                r.peak_vram_mb = meter.peak_vram_mb or np.nan
                r.roc_auc, r.log_loss, r.accuracy, r.f1_macro = evaluate(
                    yte, proba, n_classes)
                print(f"  {m.name:<14} acc={r.accuracy:.4f} auc={r.roc_auc:.4f} "
                      f"ll={r.log_loss:.4f} f1={r.f1_macro:.4f} "
                      f"fit={r.fit_s:.2f}s infer={r.latency_ms:.3f}ms/s "
                      f"ram={r.peak_ram_mb:.0f}MB vram={r.peak_vram_mb or 0:.0f}MB")
            except Exception as e:
                r.status = f"FAILED"
                r.notes = str(e).split("\n")[0][:120]
                print(f"  {m.name:<14} FAILED: {r.notes}")
            results.append(r)
    return results


# --------------------------------------------------------------------------
# Markdown model-card snippet
# --------------------------------------------------------------------------
def _bold_best(df: pd.DataFrame, col: str, higher=True) -> pd.Series:
    best = df[col].max() if higher else df[col].min()
    return df[col].map(lambda v: f"**{v:.4f}**" if v == best else f"{v:.4f}")


def to_markdown(results: list[Result]) -> str:
    df = pd.DataFrame([vars(r) for r in results])
    ok = df[df.status == "ok"]
    models = [m for m in df.model.unique()]
    lines = ["## Benchmark results",
             "",
             "Zero-shot in-context models (FelaTab, TabFM) vs trained baselines "
             "(XGBoost, LightGBM). OpenML datasets, stratified 80/20 split, seed 42.",
             ""]

    # 1) per-dataset performance tables
    for metric, higher, title in [
        ("accuracy", True, "Accuracy"), ("roc_auc", True, "ROC-AUC"),
        ("log_loss", False, "Log Loss"), ("f1_macro", True, "F1 (macro)")]:
        piv = ok.pivot_table(index="dataset", columns="model", values=metric)
        lines.append(f"### {title} per dataset")
        lines.append("")
        lines.append("| Dataset | " + " | ".join(piv.columns) + " |")
        lines.append("|" + "---|" * (len(piv.columns) + 1))
        for d, row in piv.iterrows():
            best = row.max() if higher else row.min()
            cells = [f"**{v:.4f}**" if v == best else (f"{v:.4f}" if pd.notna(v) else "—")
                     for v in row]
            lines.append(f"| {d} | " + " | ".join(cells) + " |")
        lines.append("")

    # 2) summary with average rank
    lines += ["### Summary (mean across datasets)", ""]
    summ = ok.groupby("model").agg(
        mean_acc=("accuracy", "mean"), mean_auc=("roc_auc", "mean"),
        mean_logloss=("log_loss", "mean"), mean_f1=("f1_macro", "mean"))
    ranks = []
    for metric, higher in [("accuracy", True), ("roc_auc", True),
                           ("log_loss", False), ("f1_macro", True)]:
        piv = ok.pivot_table(index="dataset", columns="model", values=metric)
        r = piv.rank(axis=1, ascending=not higher).mean()
        ranks.append(r)
    avg_rank = pd.concat(ranks, axis=1).mean(axis=1)
    summ["avg_rank"] = avg_rank
    lines.append("| Model | Mean Acc | Mean ROC-AUC | Mean LogLoss | Mean F1 | Avg Rank |")
    lines.append("|---|---|---|---|---|---|")
    for m in summ.index:
        s = summ.loc[m]
        lines.append(f"| {m} | {s.mean_acc:.4f} | {s.mean_auc:.4f} | "
                     f"{s.mean_logloss:.4f} | {s.mean_f1:.4f} | **{s.avg_rank:.2f}** |")
    lines.append("")

    # 3) efficiency
    lines += ["### Efficiency", "",
              "| Model | Device | Fit time (s) | Latency (ms/sample) | Peak RAM (MB) | Peak VRAM (MB) |",
              "|---|---|---|---|---|---|"]
    for m in models:
        sub = ok[ok.model == m]
        if sub.empty:
            continue
        dev = sub.device.iloc[0]
        vram = sub.peak_vram_mb.mean()
        lines.append(f"| {m} | {dev} | {sub.fit_s.mean():.2f} | "
                     f"{sub.latency_ms.mean():.3f} | {sub.peak_ram_mb.mean():.0f} | "
                     f"{'N/A' if pd.isna(vram) or vram == 0 else f'{vram:.0f}'} |")
    lines.append("")

    failed = df[df.status != "ok"]
    if not failed.empty:
        lines += ["<details><summary>Failed runs</summary>", "",
                  "| Dataset | Model | Reason |", "|---|---|---|"]
        for _, f in failed.iterrows():
            lines.append(f"| {f.dataset} | {f.model} | {f.notes} |")
        lines += ["", "</details>", ""]
    return "\n".join(lines)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--smoke", action="store_true")
    ap.add_argument("--device", choices=["gpu", "cpu"], default="gpu")
    ap.add_argument("--tiers", default="big,small")
    ap.add_argument("--skip", default="")
    ap.add_argument("--datasets", default="")
    args = ap.parse_args()

    results = run(args)
    df = pd.DataFrame([vars(r) for r in results])
    out_csv = HERE / "results.csv"
    df.to_csv(out_csv, index=False)
    md = to_markdown(results)
    (HERE / "model_card_snippet.md").write_text(md)
    print(f"\nwrote {out_csv} and {HERE/'model_card_snippet.md'}")
    print("\n" + md)


if __name__ == "__main__":
    main()