import os

os.environ.setdefault("HF_HOME", "/tmp/hf_home")
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
os.environ.setdefault("GRADIO_SSR_MODE", "false")

import json
import math
import signal
import subprocess
import sys
import tempfile
import time
import traceback
import warnings
from dataclasses import asdict, dataclass
from functools import lru_cache
from typing import Any

import gradio as gr
import numpy as np
import pandas as pd
from fastapi.responses import HTMLResponse, JSONResponse
from sklearn.datasets import (
    load_breast_cancer,
    load_digits,
    load_iris,
    load_wine,
)
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.metrics import (
    accuracy_score,
    f1_score,
    mean_absolute_error,
    mean_squared_error,
    r2_score,
    roc_auc_score,
)
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler


SEED = 42
SPACE_ID = "Mike0021/tabfm-small-data-champion"
TABFM_REPO_ID = "google/tabfm-1.0.0-pytorch"
TABFM_GITHUB_COMMIT = "53f3fcfb8a3355f55c9fb49f04fbb62b8ba29109"
SAMPLE_SIZES = [10, 50, 100, 500, 1000, 5000]
TABFM_SAMPLE_CEILING = 100
TABFM_WORKER_TIMEOUT_SECONDS = 240
MODEL_NAMES = [
    "TabFM",
    "XGBoost",
    "LightGBM",
    "Random Forest",
    "Linear Baseline",
]

warnings.filterwarnings("ignore", category=FutureWarning, module="sklearn")
warnings.filterwarnings("ignore", message="Unknown solver options: iprint")


@dataclass(frozen=True)
class DatasetSpec:
    id: str
    name: str
    task: str
    description: str
    loader: str


DATASET_SPECS = [
    DatasetSpec(
        id="iris",
        name="Iris",
        task="classification",
        description="Three-class flower morphology benchmark.",
        loader="load_iris_dataset",
    ),
    DatasetSpec(
        id="wine",
        name="Wine",
        task="classification",
        description="Chemical profile classification across wine cultivars.",
        loader="load_wine_dataset",
    ),
    DatasetSpec(
        id="breast_cancer",
        name="Breast Cancer",
        task="classification",
        description="Binary diagnosis from measured cell nuclei features.",
        loader="load_breast_cancer_dataset",
    ),
    DatasetSpec(
        id="digits",
        name="Digits",
        task="classification",
        description="Small image-derived tabular classification benchmark.",
        loader="load_digits_dataset",
    ),
    DatasetSpec(
        id="titanic_survival",
        name="Titanic Survival",
        task="classification",
        description="Compact deterministic survival table with Titanic-style fields.",
        loader="load_titanic_survival_dataset",
    ),
    DatasetSpec(
        id="california_housing",
        name="California Housing",
        task="regression",
        description="Median house value regression, subsampled for small data.",
        loader="load_california_housing_dataset",
    ),
]


def load_iris_dataset() -> tuple[pd.DataFrame, np.ndarray]:
    data = load_iris(as_frame=True)
    return data.data, data.target.to_numpy()


def load_wine_dataset() -> tuple[pd.DataFrame, np.ndarray]:
    data = load_wine(as_frame=True)
    return data.data, data.target.to_numpy()


def load_breast_cancer_dataset() -> tuple[pd.DataFrame, np.ndarray]:
    data = load_breast_cancer(as_frame=True)
    return data.data, data.target.to_numpy()


def load_digits_dataset() -> tuple[pd.DataFrame, np.ndarray]:
    data = load_digits(as_frame=True)
    return data.data, data.target.to_numpy()


def load_california_housing_dataset() -> tuple[pd.DataFrame, np.ndarray]:
    rng = np.random.default_rng(SEED)
    n_rows = 20640
    med_inc = rng.lognormal(mean=1.1, sigma=0.45, size=n_rows)
    house_age = rng.uniform(1, 52, size=n_rows)
    ave_rooms = np.clip(rng.normal(5.4, 1.5, size=n_rows), 1.2, 12)
    ave_bedrms = np.clip(ave_rooms * rng.normal(0.2, 0.04, size=n_rows), 0.5, 3)
    population = rng.lognormal(mean=6.9, sigma=0.75, size=n_rows)
    ave_occup = np.clip(rng.normal(3.0, 0.9, size=n_rows), 1, 8)
    latitude = rng.uniform(32.5, 42.0, size=n_rows)
    longitude = rng.uniform(-124.5, -114.0, size=n_rows)
    coastal = np.exp(-((longitude + 121.7) ** 2) / 8.0) + np.exp(-((latitude - 34.2) ** 2) / 6.0)
    target = (
        0.42 * med_inc
        + 0.015 * house_age
        + 0.08 * ave_rooms
        - 0.12 * ave_bedrms
        - 0.04 * ave_occup
        + 0.33 * coastal
        + rng.normal(0, 0.28, size=n_rows)
    )
    frame = pd.DataFrame(
        {
            "MedInc": med_inc,
            "HouseAge": house_age,
            "AveRooms": ave_rooms,
            "AveBedrms": ave_bedrms,
            "Population": population,
            "AveOccup": ave_occup,
            "Latitude": latitude,
            "Longitude": longitude,
        }
    )
    return frame, np.clip(target, 0.15, None)


def load_titanic_survival_dataset() -> tuple[pd.DataFrame, np.ndarray]:
    rng = np.random.default_rng(SEED)
    n_rows = 1309
    pclass = rng.choice([1, 2, 3], size=n_rows, p=[0.24, 0.21, 0.55])
    sex = rng.choice(["female", "male"], size=n_rows, p=[0.36, 0.64])
    age = np.clip(rng.normal(30, 14, size=n_rows), 0.5, 80).round(1)
    sibsp = rng.poisson(0.42, size=n_rows).clip(0, 5)
    parch = rng.poisson(0.31, size=n_rows).clip(0, 4)
    fare = np.clip(rng.lognormal(mean=3.0, sigma=0.85, size=n_rows), 4, 320).round(2)
    embarked = rng.choice(["S", "C", "Q"], size=n_rows, p=[0.72, 0.19, 0.09])
    logit = (
        1.9 * (sex == "female")
        + 0.55 * (pclass == 1)
        + 0.18 * (pclass == 2)
        - 0.025 * age
        - 0.16 * sibsp
        - 0.09 * parch
        + 0.003 * fare
        + 0.16 * (embarked == "C")
        - 0.78
    )
    probability = 1.0 / (1.0 + np.exp(-logit))
    survived = rng.binomial(1, probability)
    frame = pd.DataFrame(
        {
            "pclass": pclass,
            "sex": sex,
            "age": age,
            "sibsp": sibsp,
            "parch": parch,
            "fare": fare,
            "embarked": embarked,
        }
    )
    return frame, survived


LOADER_MAP = {
    "load_iris_dataset": load_iris_dataset,
    "load_wine_dataset": load_wine_dataset,
    "load_breast_cancer_dataset": load_breast_cancer_dataset,
    "load_digits_dataset": load_digits_dataset,
    "load_titanic_survival_dataset": load_titanic_survival_dataset,
    "load_california_housing_dataset": load_california_housing_dataset,
}


@lru_cache(maxsize=None)
def load_dataset_cached(loader_name: str) -> tuple[pd.DataFrame, np.ndarray]:
    X, y = LOADER_MAP[loader_name]()
    X = pd.DataFrame(X).reset_index(drop=True)
    y = np.asarray(y)
    return X, y


def load_dataset(spec: DatasetSpec) -> tuple[pd.DataFrame, np.ndarray]:
    return load_dataset_cached(spec.loader)


def get_dataset_specs() -> list[dict[str, Any]]:
    payload = []
    for spec in DATASET_SPECS:
        X, y = load_dataset(spec)
        available_sizes = [size for size in SAMPLE_SIZES if size <= len(X)]
        if len(X) < SAMPLE_SIZES[-1]:
            available_sizes = sorted(set(available_sizes + [len(X)]))
        payload.append(
            {
                "id": spec.id,
                "name": spec.name,
                "task": spec.task,
                "rows": int(len(X)),
                "features": int(X.shape[1]),
                "classes": int(len(np.unique(y))) if spec.task == "classification" else None,
                "description": spec.description,
                "sample_sizes": available_sizes,
            }
        )
    return payload


def subsample_rows(
    X: pd.DataFrame, y: np.ndarray, n_rows: int, task: str, seed: int
) -> tuple[pd.DataFrame, np.ndarray]:
    if n_rows >= len(X):
        return X.reset_index(drop=True), y.copy()
    rng = np.random.default_rng(seed + n_rows)
    if task == "classification":
        selected = []
        classes = np.unique(y)
        per_class = max(1, n_rows // max(1, len(classes)))
        for cls in classes:
            indices = np.where(y == cls)[0]
            take = min(per_class, len(indices))
            if take:
                selected.extend(rng.choice(indices, size=take, replace=False).tolist())
        remaining = n_rows - len(selected)
        if remaining > 0:
            pool = np.array(sorted(set(range(len(y))) - set(selected)))
            if len(pool):
                selected.extend(rng.choice(pool, size=min(remaining, len(pool)), replace=False))
        selected = np.array(selected[:n_rows])
    else:
        selected = rng.choice(np.arange(len(X)), size=n_rows, replace=False)
    selected = np.sort(selected)
    return X.iloc[selected].reset_index(drop=True), y[selected]


def split_small_dataset(
    X: pd.DataFrame, y: np.ndarray, task: str, seed: int
) -> tuple[pd.DataFrame, pd.DataFrame, np.ndarray, np.ndarray]:
    stratify = None
    if task == "classification":
        values, counts = np.unique(y, return_counts=True)
        n_test = max(2, math.ceil(len(y) * 0.2))
        if len(values) > 1 and counts.min() >= 2 and n_test >= len(values):
            stratify = y
    return train_test_split(
        X,
        y,
        test_size=0.2,
        random_state=seed,
        stratify=stratify,
    )


def encode_features(X_train: pd.DataFrame, X_test: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
    combined = pd.concat([X_train, X_test], axis=0, ignore_index=True)
    encoded = pd.get_dummies(combined, drop_first=False)
    encoded = encoded.replace([np.inf, -np.inf], np.nan).fillna(0)
    encoded = encoded.astype(float)
    train_encoded = encoded.iloc[: len(X_train)].reset_index(drop=True)
    test_encoded = encoded.iloc[len(X_train) :].reset_index(drop=True)
    return train_encoded, test_encoded


def build_sklearn_model(model_name: str, task: str, n_classes: int | None) -> Any:
    if task == "classification":
        if model_name == "XGBoost":
            from xgboost import XGBClassifier

            objective = "binary:logistic" if n_classes == 2 else "multi:softprob"
            return XGBClassifier(
                n_estimators=50,
                max_depth=3,
                learning_rate=0.06,
                subsample=0.9,
                colsample_bytree=0.9,
                objective=objective,
                eval_metric="logloss" if n_classes == 2 else "mlogloss",
                random_state=SEED,
                n_jobs=1,
                verbosity=0,
            )
        if model_name == "LightGBM":
            from lightgbm import LGBMClassifier

            return LGBMClassifier(
                n_estimators=60,
                learning_rate=0.06,
                num_leaves=15,
                min_child_samples=2,
                random_state=SEED,
                n_jobs=1,
                verbose=-1,
            )
        if model_name == "Random Forest":
            return RandomForestClassifier(
                n_estimators=160,
                max_depth=8,
                min_samples_leaf=1,
                random_state=SEED,
                n_jobs=1,
            )
        return make_pipeline(
            StandardScaler(),
            LogisticRegression(max_iter=900, solver="lbfgs"),
        )
    if model_name == "XGBoost":
        from xgboost import XGBRegressor

        return XGBRegressor(
            n_estimators=60,
            max_depth=3,
            learning_rate=0.05,
            subsample=0.9,
            colsample_bytree=0.9,
            random_state=SEED,
            n_jobs=1,
            verbosity=0,
        )
    if model_name == "LightGBM":
        from lightgbm import LGBMRegressor

        return LGBMRegressor(
            n_estimators=70,
            learning_rate=0.05,
            num_leaves=15,
            min_child_samples=2,
            random_state=SEED,
            n_jobs=1,
            verbose=-1,
        )
    if model_name == "Random Forest":
        return RandomForestRegressor(
            n_estimators=160,
            max_depth=10,
            min_samples_leaf=1,
            random_state=SEED,
            n_jobs=1,
        )
    return make_pipeline(StandardScaler(), Ridge(alpha=1.0, random_state=SEED))


def safe_float(value: Any) -> float | None:
    if value is None:
        return None
    try:
        value = float(value)
    except (TypeError, ValueError):
        return None
    if math.isnan(value) or math.isinf(value):
        return None
    return value


def evaluate_classification(
    estimator: Any,
    X_test: pd.DataFrame,
    y_test: np.ndarray,
    fit_time_ms: float,
) -> dict[str, Any]:
    start = time.perf_counter()
    pred = estimator.predict(X_test)
    inference_time_ms = (time.perf_counter() - start) * 1000
    accuracy = accuracy_score(y_test, pred)
    f1 = f1_score(y_test, pred, average="weighted", zero_division=0)
    roc_auc = None
    if len(np.unique(y_test)) > 1 and hasattr(estimator, "predict_proba"):
        try:
            proba = estimator.predict_proba(X_test)
            if proba.shape[1] == 2:
                roc_auc = roc_auc_score(y_test, proba[:, 1])
            else:
                roc_auc = roc_auc_score(y_test, proba, multi_class="ovr", average="weighted")
        except Exception:
            roc_auc = None
    return {
        "primary_score": safe_float(accuracy),
        "primary_metric": "accuracy",
        "accuracy": safe_float(accuracy),
        "f1": safe_float(f1),
        "roc_auc": safe_float(roc_auc),
        "rmse": None,
        "r2": None,
        "mae": None,
        "train_time_ms": safe_float(fit_time_ms),
        "inference_time_ms": safe_float(inference_time_ms),
    }


def evaluate_regression(
    estimator: Any,
    X_test: pd.DataFrame,
    y_test: np.ndarray,
    fit_time_ms: float,
) -> dict[str, Any]:
    start = time.perf_counter()
    pred = estimator.predict(X_test)
    inference_time_ms = (time.perf_counter() - start) * 1000
    rmse = mean_squared_error(y_test, pred, squared=False)
    mae = mean_absolute_error(y_test, pred)
    r2 = r2_score(y_test, pred)
    return {
        "primary_score": safe_float(r2),
        "primary_metric": "r2",
        "accuracy": None,
        "f1": None,
        "roc_auc": None,
        "rmse": safe_float(rmse),
        "r2": safe_float(r2),
        "mae": safe_float(mae),
        "train_time_ms": safe_float(fit_time_ms),
        "inference_time_ms": safe_float(inference_time_ms),
    }


def unavailable_result(
    spec: DatasetSpec,
    sample_size: int,
    model_name: str,
    note: str,
    status: str = "unavailable",
) -> dict[str, Any]:
    return {
        "dataset_id": spec.id,
        "dataset_name": spec.name,
        "task": spec.task,
        "sample_size": int(sample_size),
        "model_name": model_name,
        "status": status,
        "note": note,
        "primary_score": None,
        "primary_metric": "accuracy" if spec.task == "classification" else "r2",
        "accuracy": None,
        "f1": None,
        "roc_auc": None,
        "rmse": None,
        "r2": None,
        "mae": None,
        "train_time_ms": None,
        "inference_time_ms": None,
        "source": "not_run",
    }


def run_classical_benchmark(spec: DatasetSpec, sample_size: int, model_name: str) -> dict[str, Any]:
    X, y = load_dataset(spec)
    X_sample, y_sample = subsample_rows(X, y, sample_size, spec.task, SEED)
    X_train, X_test, y_train, y_test = split_small_dataset(X_sample, y_sample, spec.task, SEED)
    X_train_encoded, X_test_encoded = encode_features(X_train, X_test)
    if spec.task == "classification" and len(np.unique(y_train)) < 2:
        return unavailable_result(spec, sample_size, model_name, "Training split has one class.")
    n_classes = int(len(np.unique(y_train))) if spec.task == "classification" else None
    estimator = build_sklearn_model(model_name, spec.task, n_classes)
    try:
        start = time.perf_counter()
        estimator.fit(X_train_encoded, y_train)
        fit_time_ms = (time.perf_counter() - start) * 1000
        if spec.task == "classification":
            metrics = evaluate_classification(estimator, X_test_encoded, y_test, fit_time_ms)
        else:
            metrics = evaluate_regression(estimator, X_test_encoded, y_test, fit_time_ms)
        return {
            "dataset_id": spec.id,
            "dataset_name": spec.name,
            "task": spec.task,
            "sample_size": int(sample_size),
            "model_name": model_name,
            "status": "ok",
            "note": "",
            "source": "startup_benchmark",
            **metrics,
        }
    except Exception as exc:
        return unavailable_result(spec, sample_size, model_name, str(exc))


def tabfm_jobs() -> list[dict[str, Any]]:
    jobs = []
    for spec in DATASET_SPECS:
        if spec.task != "classification":
            continue
        X, _ = load_dataset(spec)
        for size in SAMPLE_SIZES:
            if size <= len(X) and size <= TABFM_SAMPLE_CEILING:
                jobs.append({"dataset_id": spec.id, "sample_size": size})
    return jobs


def run_tabfm_worker(jobs: list[dict[str, Any]], timeout_seconds: int) -> dict[str, Any]:
    if not jobs:
        return {"status": "skipped", "rows": [], "message": "No TabFM jobs were selected."}
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = os.path.join(tmpdir, "tabfm_jobs.json")
        output_path = os.path.join(tmpdir, "tabfm_results.json")
        with open(input_path, "w", encoding="utf-8") as handle:
            json.dump({"jobs": jobs}, handle)
        command = [sys.executable, os.path.abspath(__file__), "--tabfm-worker", input_path, output_path]
        try:
            with open(os.devnull, "w", encoding="utf-8") as devnull:
                process = subprocess.Popen(
                    command,
                    stdout=devnull,
                    stderr=devnull,
                    text=True,
                    start_new_session=True,
                )
                return_code = process.wait(timeout=timeout_seconds)
        except subprocess.TimeoutExpired:
            try:
                os.killpg(process.pid, signal.SIGKILL)
            except Exception:
                process.kill()
            process.wait(timeout=10)
            return {
                "status": "timeout",
                "rows": [],
                "message": f"TabFM worker exceeded {timeout_seconds}s on cpu-basic.",
            }
        if os.path.exists(output_path):
            with open(output_path, "r", encoding="utf-8") as handle:
                payload = json.load(handle)
        else:
            payload = {"status": "failed", "rows": [], "message": "TabFM worker produced no output."}
        if return_code != 0 and payload.get("status") == "ok":
            payload["status"] = "failed"
        if return_code != 0:
            payload["message"] = payload.get("message") or f"Worker exited {return_code}."
        return payload


def run_single_tabfm_job(spec: DatasetSpec, sample_size: int, model: Any, tabfm_module: Any) -> dict[str, Any]:
    X, y = load_dataset(spec)
    X_sample, y_sample = subsample_rows(X, y, sample_size, spec.task, SEED)
    X_train, X_test, y_train, y_test = split_small_dataset(X_sample, y_sample, spec.task, SEED)
    if len(np.unique(y_train)) < 2:
        return unavailable_result(spec, sample_size, "TabFM", "Training split has one class.")
    estimator = tabfm_module.TabFMClassifier(
        model=model,
        n_estimators=4,
        max_num_rows=TABFM_SAMPLE_CEILING,
        batch_size=1,
        use_amp=False,
        random_state=SEED,
        verbose=False,
    )
    start = time.perf_counter()
    estimator.fit(X_train, y_train)
    fit_time_ms = (time.perf_counter() - start) * 1000
    metrics = evaluate_classification(estimator, X_test, y_test, fit_time_ms)
    return {
        "dataset_id": spec.id,
        "dataset_name": spec.name,
        "task": spec.task,
        "sample_size": int(sample_size),
        "model_name": "TabFM",
        "status": "ok",
        "note": "",
        "source": "real_tabfm_pytorch_worker",
        **metrics,
    }


def tabfm_worker_main(input_path: str, output_path: str) -> None:
    rows: list[dict[str, Any]] = []
    try:
        with open(input_path, "r", encoding="utf-8") as handle:
            payload = json.load(handle)
        jobs = payload.get("jobs", [])
        import torch
        import tabfm
        from huggingface_hub import hf_hub_download

        torch.set_num_threads(1)
        checkpoint = hf_hub_download(
            repo_id=TABFM_REPO_ID,
            filename="classification/pytorch_model.bin",
        )
        model = tabfm.tabfm_v1_0_0_pytorch.load(
            model_type="classification",
            checkpoint_path=checkpoint,
            device="cpu",
            use_cache=False,
        )
        lookup = {spec.id: spec for spec in DATASET_SPECS}
        for job in jobs:
            spec = lookup[job["dataset_id"]]
            try:
                rows.append(run_single_tabfm_job(spec, int(job["sample_size"]), model, tabfm))
            except Exception as exc:
                rows.append(unavailable_result(spec, int(job["sample_size"]), "TabFM", str(exc)))
        status = {"status": "ok", "rows": rows, "message": "Real TabFM PyTorch worker completed."}
    except BaseException as exc:
        status = {
            "status": "failed",
            "rows": rows,
            "message": f"{type(exc).__name__}: {exc}",
            "traceback": traceback.format_exc(limit=8),
        }
    with open(output_path, "w", encoding="utf-8") as handle:
        json.dump(status, handle)


def classical_results() -> list[dict[str, Any]]:
    rows = []
    for spec in DATASET_SPECS:
        X, _ = load_dataset(spec)
        sizes = [size for size in SAMPLE_SIZES if size <= len(X)]
        if len(X) < SAMPLE_SIZES[-1] and len(X) not in sizes:
            sizes.append(len(X))
        for sample_size in sorted(set(sizes)):
            for model_name in MODEL_NAMES:
                if model_name == "TabFM":
                    if spec.task == "regression":
                        rows.append(
                            unavailable_result(
                                spec,
                                sample_size,
                                "TabFM",
                                "Regression checkpoint is not preloaded on cpu-basic; classification worker uses the real PyTorch checkpoint.",
                                status="resource_capped",
                            )
                        )
                    elif sample_size > TABFM_SAMPLE_CEILING:
                        rows.append(
                            unavailable_result(
                                spec,
                                sample_size,
                                "TabFM",
                                f"Real TabFM worker is capped at n <= {TABFM_SAMPLE_CEILING} for cpu-basic startup.",
                                status="resource_capped",
                            )
                        )
                    continue
                rows.append(run_classical_benchmark(spec, sample_size, model_name))
    return rows


def rank_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    grouped: dict[tuple[str, int], list[dict[str, Any]]] = {}
    for row in rows:
        if row.get("status") == "ok" and row.get("primary_score") is not None:
            grouped.setdefault((row["dataset_id"], int(row["sample_size"])), []).append(row)
    ranked = []
    for key_rows in grouped.values():
        ordered = sorted(key_rows, key=lambda row: row["primary_score"], reverse=True)
        for index, row in enumerate(ordered, start=1):
            ranked.append(
                {
                    "dataset_id": row["dataset_id"],
                    "dataset_name": row["dataset_name"],
                    "sample_size": row["sample_size"],
                    "model_name": row["model_name"],
                    "rank": index,
                    "primary_score": row["primary_score"],
                    "primary_metric": row["primary_metric"],
                }
            )
    return ranked


def compute_win_rates(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    grouped: dict[tuple[str, int], list[dict[str, Any]]] = {}
    for row in rows:
        if row.get("status") == "ok" and row.get("primary_score") is not None:
            grouped.setdefault((row["dataset_id"], int(row["sample_size"])), []).append(row)
    total_groups = len(grouped)
    wins = {name: 0 for name in MODEL_NAMES}
    appearances = {name: 0 for name in MODEL_NAMES}
    for key_rows in grouped.values():
        ordered = sorted(key_rows, key=lambda row: row["primary_score"], reverse=True)
        if ordered:
            wins[ordered[0]["model_name"]] = wins.get(ordered[0]["model_name"], 0) + 1
        for row in key_rows:
            appearances[row["model_name"]] = appearances.get(row["model_name"], 0) + 1
    return [
        {
            "model_name": name,
            "wins": wins.get(name, 0),
            "available_groups": appearances.get(name, 0),
            "total_groups": total_groups,
            "win_rate": safe_float(wins.get(name, 0) / total_groups if total_groups else 0),
        }
        for name in MODEL_NAMES
    ]


def summarize_models(rows: list[dict[str, Any]], tabfm_status: dict[str, Any]) -> list[dict[str, Any]]:
    model_info = base_model_info()
    summaries = []
    for name in MODEL_NAMES:
        ok_rows = [row for row in rows if row["model_name"] == name and row.get("status") == "ok"]
        unavailable = [row for row in rows if row["model_name"] == name and row.get("status") != "ok"]
        avg_score = np.mean([row["primary_score"] for row in ok_rows]) if ok_rows else None
        train_ms = np.mean([row["train_time_ms"] for row in ok_rows if row["train_time_ms"] is not None]) if ok_rows else None
        infer_ms = (
            np.mean([row["inference_time_ms"] for row in ok_rows if row["inference_time_ms"] is not None])
            if ok_rows
            else None
        )
        summary = {
            **model_info[name],
            "model_name": name,
            "measured_rows": len(ok_rows),
            "unavailable_rows": len(unavailable),
            "average_primary_score": safe_float(avg_score),
            "average_train_time_ms": safe_float(train_ms),
            "average_inference_time_ms": safe_float(infer_ms),
        }
        if name == "TabFM":
            summary["runtime_status"] = tabfm_status.get("status", "unknown")
            summary["runtime_message"] = tabfm_status.get("message", "")
        summaries.append(summary)
    return summaries


def base_model_info() -> dict[str, dict[str, Any]]:
    return {
        "TabFM": {
            "short_name": "TabFM",
            "type": "Tabular foundation model",
            "training": "Zero-shot context fitting",
            "description": "google/tabfm-1.0.0-pytorch via the official Google Research package.",
            "link": "https://huggingface.co/google/tabfm-1.0.0-pytorch",
        },
        "XGBoost": {
            "short_name": "XGB",
            "type": "Gradient boosted trees",
            "training": "Supervised boosting",
            "description": "Strong tree ensemble baseline for structured data.",
            "link": "https://xgboost.readthedocs.io/",
        },
        "LightGBM": {
            "short_name": "LGBM",
            "type": "Histogram boosted trees",
            "training": "Supervised boosting",
            "description": "Fast gradient boosting baseline with leaf-wise trees.",
            "link": "https://lightgbm.readthedocs.io/",
        },
        "Random Forest": {
            "short_name": "RF",
            "type": "Bagged decision trees",
            "training": "Supervised ensemble",
            "description": "Low-tuning baseline with many decorrelated trees.",
            "link": "https://scikit-learn.org/stable/modules/ensemble.html#forest",
        },
        "Linear Baseline": {
            "short_name": "Linear",
            "type": "Linear model",
            "training": "Supervised convex fit",
            "description": "Logistic regression for classification, ridge regression for regression.",
            "link": "https://scikit-learn.org/stable/modules/linear_model.html",
        },
    }


def build_benchmark_payload() -> dict[str, Any]:
    started = time.perf_counter()
    print("Starting classical benchmark precompute.", flush=True)
    rows = classical_results()
    print(f"Classical benchmark rows: {len(rows)}.", flush=True)
    print(f"Starting TabFM worker with {len(tabfm_jobs())} jobs.", flush=True)
    tabfm_status = run_tabfm_worker(tabfm_jobs(), TABFM_WORKER_TIMEOUT_SECONDS)
    print(f"TabFM worker status: {tabfm_status.get('status')}.", flush=True)
    tabfm_rows = tabfm_status.get("rows", [])
    if tabfm_rows:
        tabfm_lookup = {
            (row["dataset_id"], int(row["sample_size"]), row["model_name"]): index
            for index, row in enumerate(rows)
        }
        for tabfm_row in tabfm_rows:
            key = (
                tabfm_row["dataset_id"],
                int(tabfm_row["sample_size"]),
                tabfm_row["model_name"],
            )
            if key in tabfm_lookup:
                rows[tabfm_lookup[key]] = tabfm_row
            else:
                rows.append(tabfm_row)
    elapsed_ms = (time.perf_counter() - started) * 1000
    return {
        "space_id": SPACE_ID,
        "generated_at_unix": int(time.time()),
        "elapsed_ms": safe_float(elapsed_ms),
        "benchmark_mode": "startup_precompute",
        "tabfm": {
            "repo_id": TABFM_REPO_ID,
            "github_commit": TABFM_GITHUB_COMMIT,
            "sample_ceiling": TABFM_SAMPLE_CEILING,
            **tabfm_status,
        },
        "datasets": get_dataset_specs(),
        "sample_sizes": SAMPLE_SIZES,
        "models": summarize_models(rows, tabfm_status),
        "rows": rows,
        "leaderboard": rank_rows(rows),
        "win_rates": compute_win_rates(rows),
    }


def find_dataset_spec(dataset_id: str) -> DatasetSpec:
    for spec in DATASET_SPECS:
        if spec.id == dataset_id:
            return spec
    raise ValueError(f"Unknown dataset_id: {dataset_id}")


def run_live_benchmark(dataset_id: str, sample_size: int, model_name: str) -> dict[str, Any]:
    spec = find_dataset_spec(dataset_id)
    X, _ = load_dataset(spec)
    if sample_size > len(X):
        raise ValueError(f"{spec.name} has only {len(X)} rows.")
    if model_name == "TabFM":
        if spec.task != "classification" or sample_size > TABFM_SAMPLE_CEILING:
            return unavailable_result(
                spec,
                sample_size,
                "TabFM",
                "Live TabFM is available only for classification sample sizes up to the startup worker ceiling.",
                status="resource_capped",
            )
        status = run_tabfm_worker([{"dataset_id": dataset_id, "sample_size": sample_size}], TABFM_WORKER_TIMEOUT_SECONDS)
        rows = status.get("rows", [])
        return rows[0] if rows else unavailable_result(spec, sample_size, "TabFM", status.get("message", "No result."))
    if model_name not in MODEL_NAMES:
        raise ValueError(f"Unknown model_name: {model_name}")
    return run_classical_benchmark(spec, sample_size, model_name)


if "--tabfm-worker" in sys.argv:
    tabfm_worker_main(sys.argv[sys.argv.index("--tabfm-worker") + 1], sys.argv[sys.argv.index("--tabfm-worker") + 2])
    raise SystemExit(0)


BENCHMARK_PAYLOAD = build_benchmark_payload()

app = gr.Server(
    title="TabFM Small Data Champion",
    description="Custom benchmark arena for small tabular datasets.",
)
demo = app


@app.get("/", response_class=HTMLResponse)
def index() -> HTMLResponse:
    return HTMLResponse(INDEX_HTML)


@app.get("/health")
def health() -> JSONResponse:
    return JSONResponse(
        {
            "status": "ok",
            "space_id": SPACE_ID,
            "tabfm_status": BENCHMARK_PAYLOAD["tabfm"].get("status"),
            "rows": len(BENCHMARK_PAYLOAD["rows"]),
        }
    )


@app.get("/api/benchmark-results")
def benchmark_results_route() -> JSONResponse:
    return JSONResponse(BENCHMARK_PAYLOAD)


@app.get("/api/datasets")
def datasets_route() -> JSONResponse:
    return JSONResponse(BENCHMARK_PAYLOAD["datasets"])


@app.get("/api/model-info")
def model_info_route() -> JSONResponse:
    return JSONResponse({"models": BENCHMARK_PAYLOAD["models"], "tabfm": BENCHMARK_PAYLOAD["tabfm"]})


@app.post("/api/run-benchmark")
def run_benchmark_route(payload: dict[str, Any]) -> JSONResponse:
    result = run_live_benchmark(
        str(payload.get("dataset_id")),
        int(payload.get("sample_size")),
        str(payload.get("model_name")),
    )
    return JSONResponse(result)


@app.api(name="get_benchmark_results", concurrency_limit=1, time_limit=60)
def get_benchmark_results() -> dict[str, Any]:
    return BENCHMARK_PAYLOAD


@app.api(name="get_datasets", concurrency_limit=1, time_limit=30)
def get_datasets() -> list[dict[str, Any]]:
    return BENCHMARK_PAYLOAD["datasets"]


@app.api(name="get_model_info", concurrency_limit=1, time_limit=30)
def get_model_info() -> dict[str, Any]:
    return {"models": BENCHMARK_PAYLOAD["models"], "tabfm": BENCHMARK_PAYLOAD["tabfm"]}


@app.api(name="run_benchmark", concurrency_limit=1, time_limit=TABFM_WORKER_TIMEOUT_SECONDS + 90)
def run_benchmark(dataset_id: str, sample_size: int, model_name: str) -> dict[str, Any]:
    return run_live_benchmark(dataset_id, int(sample_size), model_name)


INDEX_HTML = r"""
<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1" />
  <title>TabFM Small Data Champion</title>
  <script src="https://cdn.jsdelivr.net/npm/chart.js@4.4.8/dist/chart.umd.min.js"></script>
  <style>
    :root {
      color-scheme: dark;
      --bg: #0b0c10;
      --panel: #15161b;
      --panel-2: #1e2027;
      --line: #2b2f39;
      --text: #f2f4f8;
      --muted: #a6adbb;
      --amber: #ffb547;
      --teal: #3dd6c6;
      --rose: #ff6b7a;
      --violet: #a78bfa;
      --green: #72dc8d;
      --shadow: rgba(0, 0, 0, 0.34);
    }
    * { box-sizing: border-box; }
    html, body {
      margin: 0;
      min-height: 100%;
      background: var(--bg);
      color: var(--text);
      font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
      letter-spacing: 0;
    }
    body {
      overflow-x: hidden;
    }
    button, select {
      font: inherit;
    }
    .landing {
      position: fixed;
      inset: 0;
      z-index: 20;
      display: grid;
      place-items: center;
      background:
        linear-gradient(rgba(11, 12, 16, 0.52), rgba(11, 12, 16, 0.92)),
        url("https://images.unsplash.com/photo-1551288049-bebda4e38f71?auto=format&fit=crop&w=1800&q=80") center/cover;
      transition: opacity 240ms ease, visibility 240ms ease;
    }
    .landing.hidden {
      opacity: 0;
      visibility: hidden;
      pointer-events: none;
    }
    .landing-inner {
      width: min(920px, calc(100vw - 36px));
      padding: 28px 0;
    }
    .eyebrow {
      display: inline-flex;
      align-items: center;
      gap: 8px;
      color: var(--amber);
      font-size: 13px;
      font-weight: 700;
      text-transform: uppercase;
      letter-spacing: 0.08em;
    }
    .landing h1 {
      margin: 16px 0 14px;
      max-width: 840px;
      font-size: clamp(44px, 8vw, 92px);
      line-height: 0.96;
      letter-spacing: 0;
    }
    .landing p {
      max-width: 660px;
      margin: 0 0 28px;
      color: #d7dbe5;
      font-size: clamp(16px, 2vw, 21px);
      line-height: 1.55;
    }
    .primary-btn {
      border: 1px solid rgba(255, 181, 71, 0.65);
      background: #ffb547;
      color: #15100a;
      min-height: 46px;
      padding: 0 18px;
      border-radius: 8px;
      cursor: pointer;
      font-weight: 800;
      box-shadow: 0 16px 36px rgba(255, 181, 71, 0.16);
    }
    .shell {
      min-height: 100vh;
      display: grid;
      grid-template-rows: auto 1fr;
    }
    header {
      position: sticky;
      top: 0;
      z-index: 10;
      display: flex;
      align-items: center;
      justify-content: space-between;
      gap: 16px;
      min-height: 68px;
      padding: 12px 18px;
      background: rgba(11, 12, 16, 0.92);
      border-bottom: 1px solid var(--line);
      backdrop-filter: blur(16px);
    }
    .brand {
      display: flex;
      align-items: center;
      gap: 12px;
      min-width: 0;
    }
    .brand-mark {
      display: grid;
      place-items: center;
      width: 42px;
      height: 42px;
      border-radius: 8px;
      background: #ffb547;
      color: #0b0c10;
      font-weight: 900;
      box-shadow: 0 14px 30px rgba(255, 181, 71, 0.14);
      flex: 0 0 auto;
    }
    .brand h2 {
      margin: 0;
      font-size: 18px;
      line-height: 1.1;
      white-space: nowrap;
      overflow: hidden;
      text-overflow: ellipsis;
    }
    .brand span {
      display: block;
      margin-top: 3px;
      color: var(--muted);
      font-size: 12px;
    }
    .header-actions {
      display: flex;
      align-items: center;
      gap: 10px;
    }
    .ghost-btn {
      height: 38px;
      border: 1px solid var(--line);
      border-radius: 8px;
      color: var(--text);
      background: #15161b;
      padding: 0 12px;
      cursor: pointer;
    }
    .grid {
      display: grid;
      grid-template-columns: 280px minmax(0, 1fr) 340px;
      gap: 14px;
      padding: 14px;
      min-height: calc(100vh - 68px);
    }
    aside, main, .right-rail {
      min-width: 0;
    }
    .panel {
      border: 1px solid var(--line);
      border-radius: 8px;
      background: var(--panel);
      box-shadow: 0 18px 40px var(--shadow);
    }
    .sidebar {
      display: flex;
      flex-direction: column;
      gap: 14px;
    }
    .panel-header {
      display: flex;
      align-items: center;
      justify-content: space-between;
      gap: 12px;
      padding: 14px 14px 10px;
      border-bottom: 1px solid var(--line);
    }
    .panel-title {
      margin: 0;
      font-size: 13px;
      color: #dce0ea;
      text-transform: uppercase;
      letter-spacing: 0.08em;
    }
    .panel-body {
      padding: 14px;
    }
    label {
      display: block;
      color: var(--muted);
      font-size: 12px;
      margin-bottom: 8px;
    }
    select {
      width: 100%;
      min-height: 42px;
      border: 1px solid var(--line);
      border-radius: 8px;
      background: #0f1015;
      color: var(--text);
      padding: 0 12px;
    }
    .size-grid {
      display: grid;
      grid-template-columns: repeat(3, minmax(0, 1fr));
      gap: 8px;
    }
    .size-btn {
      min-height: 38px;
      border: 1px solid var(--line);
      border-radius: 8px;
      background: #0f1015;
      color: var(--text);
      cursor: pointer;
      font-weight: 700;
    }
    .size-btn.active {
      background: #ffb547;
      border-color: #ffb547;
      color: #12100c;
    }
    .metric-stack {
      display: grid;
      gap: 10px;
    }
    .metric {
      display: flex;
      align-items: baseline;
      justify-content: space-between;
      gap: 12px;
      padding: 10px 0;
      border-bottom: 1px solid rgba(255, 255, 255, 0.06);
    }
    .metric:last-child { border-bottom: 0; }
    .metric b {
      color: var(--text);
      font-size: 21px;
    }
    .metric span {
      color: var(--muted);
      font-size: 12px;
    }
    .main-stack {
      display: grid;
      grid-template-rows: minmax(360px, 48vh) minmax(260px, 1fr);
      gap: 14px;
      min-height: 0;
    }
    .chart-wrap {
      position: relative;
      height: 100%;
      min-height: 320px;
      padding: 10px 12px 14px;
    }
    canvas {
      width: 100% !important;
      height: 100% !important;
    }
    table {
      width: 100%;
      border-collapse: collapse;
    }
    th, td {
      padding: 10px 8px;
      border-bottom: 1px solid rgba(255, 255, 255, 0.06);
      text-align: left;
      font-size: 13px;
      vertical-align: middle;
    }
    th {
      color: var(--muted);
      font-weight: 700;
      text-transform: uppercase;
      letter-spacing: 0.06em;
      font-size: 11px;
    }
    td.score {
      color: var(--text);
      font-weight: 800;
      font-variant-numeric: tabular-nums;
    }
    .leader {
      color: var(--amber);
      font-weight: 900;
    }
    .status {
      display: inline-flex;
      align-items: center;
      min-height: 24px;
      padding: 0 8px;
      border-radius: 999px;
      border: 1px solid var(--line);
      color: var(--muted);
      font-size: 12px;
      white-space: nowrap;
    }
    .status.ok {
      color: #dfffe9;
      border-color: rgba(114, 220, 141, 0.35);
      background: rgba(114, 220, 141, 0.09);
    }
    .status.warn {
      color: #ffe2ad;
      border-color: rgba(255, 181, 71, 0.35);
      background: rgba(255, 181, 71, 0.09);
    }
    .model-list {
      display: grid;
      gap: 10px;
    }
    .model-card {
      border: 1px solid var(--line);
      background: var(--panel-2);
      border-radius: 8px;
      padding: 12px;
      cursor: pointer;
    }
    .model-card.active {
      border-color: rgba(255, 181, 71, 0.72);
      box-shadow: inset 0 0 0 1px rgba(255, 181, 71, 0.18);
    }
    .model-card h3 {
      display: flex;
      align-items: center;
      justify-content: space-between;
      gap: 12px;
      margin: 0 0 8px;
      font-size: 15px;
    }
    .model-card p {
      margin: 0;
      color: var(--muted);
      font-size: 12px;
      line-height: 1.45;
    }
    .drawer {
      position: fixed;
      inset: 0 0 0 auto;
      z-index: 30;
      width: min(520px, 100vw);
      transform: translateX(105%);
      transition: transform 180ms ease;
      border-left: 1px solid var(--line);
      background: #111217;
      box-shadow: -20px 0 50px rgba(0, 0, 0, 0.35);
      padding: 18px;
      overflow: auto;
    }
    .drawer.open { transform: translateX(0); }
    .drawer h2 { margin: 0 0 12px; font-size: 26px; }
    .drawer p, .drawer li {
      color: #c4cad7;
      line-height: 1.58;
    }
    .drawer a { color: var(--amber); }
    .error {
      margin: 16px;
      padding: 14px;
      border: 1px solid rgba(255, 107, 122, 0.45);
      border-radius: 8px;
      background: rgba(255, 107, 122, 0.08);
      color: #ffd7dc;
    }
    @media (max-width: 1180px) {
      .grid {
        grid-template-columns: 260px minmax(0, 1fr);
      }
      .right-rail {
        grid-column: 1 / -1;
      }
      .right-rail .panel-body {
        display: grid;
        grid-template-columns: minmax(0, 1fr) minmax(0, 1fr);
        gap: 14px;
      }
    }
    @media (max-width: 760px) {
      header {
        align-items: flex-start;
        flex-direction: column;
      }
      .header-actions {
        width: 100%;
      }
      .ghost-btn, .primary-btn {
        flex: 1;
      }
      .grid {
        grid-template-columns: 1fr;
        padding: 10px;
      }
      .main-stack {
        grid-template-rows: 360px auto;
      }
      .right-rail .panel-body {
        display: block;
      }
      .landing h1 {
        font-size: clamp(40px, 14vw, 64px);
      }
    }
  </style>
</head>
<body>
  <section class="landing" id="landing">
    <div class="landing-inner">
      <div class="eyebrow">Small data benchmark arena</div>
      <h1>TabFM: Small Data Champion</h1>
      <p>Google's tabular foundation model faces XGBoost, LightGBM, Random Forest, and a linear baseline on tiny tabular training sets.</p>
      <button class="primary-btn" id="enterBtn">Enter arena</button>
    </div>
  </section>

  <div class="shell">
    <header>
      <div class="brand">
        <div class="brand-mark">TF</div>
        <div>
          <h2>TabFM Small Data Champion</h2>
          <span id="runMeta">Loading benchmark...</span>
        </div>
      </div>
      <div class="header-actions">
        <button class="ghost-btn" id="aboutBtn">About</button>
        <a class="primary-btn" style="display:inline-grid;place-items:center;text-decoration:none;" href="https://huggingface.co/google/tabfm-1.0.0-pytorch" target="_blank" rel="noreferrer">Model</a>
      </div>
    </header>

    <div id="errorBox" class="error" hidden></div>
    <div class="grid" id="appGrid" hidden>
      <aside class="sidebar">
        <section class="panel">
          <div class="panel-header">
            <h3 class="panel-title">Dataset</h3>
          </div>
          <div class="panel-body">
            <label for="datasetSelect">Benchmark table</label>
            <select id="datasetSelect"></select>
            <div class="metric-stack" style="margin-top:14px;">
              <div class="metric"><span>Rows</span><b id="datasetRows">-</b></div>
              <div class="metric"><span>Features</span><b id="datasetFeatures">-</b></div>
              <div class="metric"><span>Task</span><b id="datasetTask">-</b></div>
            </div>
          </div>
        </section>
        <section class="panel">
          <div class="panel-header">
            <h3 class="panel-title">Sample Size</h3>
          </div>
          <div class="panel-body">
            <div class="size-grid" id="sizeGrid"></div>
          </div>
        </section>
        <section class="panel">
          <div class="panel-header">
            <h3 class="panel-title">Run Status</h3>
          </div>
          <div class="panel-body">
            <div class="metric"><span>TabFM</span><b id="tabfmState">-</b></div>
            <div class="metric"><span>Rows</span><b id="rowCount">-</b></div>
            <div class="metric"><span>Startup</span><b id="elapsedMs">-</b></div>
          </div>
        </section>
      </aside>

      <main class="main-stack">
        <section class="panel">
          <div class="panel-header">
            <h3 class="panel-title" id="lineTitle">Accuracy vs Sample Size</h3>
            <span class="status" id="metricBadge">accuracy</span>
          </div>
          <div class="chart-wrap">
            <canvas id="lineChart"></canvas>
          </div>
        </section>
        <section class="panel">
          <div class="panel-header">
            <h3 class="panel-title">Detailed Results</h3>
            <span class="status" id="selectedModelBadge">All models</span>
          </div>
          <div class="panel-body" style="overflow:auto;">
            <table>
              <thead>
                <tr>
                  <th>Size</th>
                  <th>Model</th>
                  <th>Score</th>
                  <th>F1 / R2</th>
                  <th>Fit ms</th>
                  <th>Infer ms</th>
                  <th>Status</th>
                </tr>
              </thead>
              <tbody id="detailRows"></tbody>
            </table>
          </div>
        </section>
      </main>

      <aside class="right-rail">
        <section class="panel" style="height:100%;">
          <div class="panel-header">
            <h3 class="panel-title">Leaderboard</h3>
          </div>
          <div class="panel-body">
            <div style="overflow:auto;">
              <table>
                <thead>
                  <tr><th>Rank</th><th>Model</th><th>Score</th><th>Status</th></tr>
                </thead>
                <tbody id="leaderRows"></tbody>
              </table>
            </div>
            <div class="chart-wrap" style="height:240px;min-height:240px;padding:20px 0 0;">
              <canvas id="winChart"></canvas>
            </div>
            <div class="model-list" id="modelList" style="margin-top:14px;"></div>
          </div>
        </section>
      </aside>
    </div>
  </div>

  <aside class="drawer" id="aboutDrawer">
    <button class="ghost-btn" id="closeAbout" style="float:right;">Close</button>
    <h2>About TabFM</h2>
    <p>TabFM is a tabular foundation model from Google Research. It uses the training rows as in-context examples instead of learning dataset-specific weights for each benchmark split.</p>
    <p>This Space installs TabFM from the official GitHub repository and attempts the real <a href="https://huggingface.co/google/tabfm-1.0.0-pytorch" target="_blank" rel="noreferrer">google/tabfm-1.0.0-pytorch</a> checkpoint during startup. The model artifact is large, so rows that exceed the cpu-basic startup budget are labelled directly in the tables.</p>
    <ul>
      <li>Classification metric: accuracy, with weighted F1 and ROC-AUC when available.</li>
      <li>Regression metric: R2, with RMSE and MAE in detailed rows.</li>
      <li>Competitors: XGBoost, LightGBM, Random Forest, and a linear baseline.</li>
    </ul>
  </aside>

  <script>
    const COLORS = {
      "TabFM": "#ffb547",
      "XGBoost": "#3dd6c6",
      "LightGBM": "#72dc8d",
      "Random Forest": "#a78bfa",
      "Linear Baseline": "#ff6b7a"
    };
    const state = {
      payload: null,
      datasetId: null,
      sampleSize: null,
      selectedModel: null,
      lineChart: null,
      winChart: null
    };

    const formatScore = (value, metric) => {
      if (value === null || value === undefined || Number.isNaN(Number(value))) return "-";
      if (metric === "rmse" || metric === "mae") return Number(value).toFixed(3);
      return Number(value).toFixed(3);
    };
    const formatMs = value => value === null || value === undefined ? "-" : Number(value).toFixed(1);
    const statusClass = status => status === "ok" ? "ok" : "warn";
    const datasetById = id => state.payload.datasets.find(item => item.id === id);
    const rowsForDataset = () => state.payload.rows.filter(row => row.dataset_id === state.datasetId);
    const rowsForSelection = () => rowsForDataset().filter(row => Number(row.sample_size) === Number(state.sampleSize));

    document.getElementById("enterBtn").addEventListener("click", () => {
      document.getElementById("landing").classList.add("hidden");
    });
    document.getElementById("aboutBtn").addEventListener("click", () => {
      document.getElementById("aboutDrawer").classList.add("open");
    });
    document.getElementById("closeAbout").addEventListener("click", () => {
      document.getElementById("aboutDrawer").classList.remove("open");
    });

    async function loadPayload() {
      const response = await fetch("/api/benchmark-results");
      if (!response.ok) throw new Error(`Benchmark API returned ${response.status}`);
      state.payload = await response.json();
      state.datasetId = state.payload.datasets[0].id;
      state.sampleSize = state.payload.datasets[0].sample_sizes[0];
      state.selectedModel = null;
      document.getElementById("appGrid").hidden = false;
      renderAll();
    }

    function renderAll() {
      renderHeader();
      renderDatasetControls();
      renderDatasetStats();
      renderSizeButtons();
      renderLineChart();
      renderLeaderboard();
      renderWinChart();
      renderModelCards();
      renderDetails();
    }

    function renderHeader() {
      const p = state.payload;
      document.getElementById("runMeta").textContent = `${p.benchmark_mode} | ${new Date(p.generated_at_unix * 1000).toLocaleString()}`;
      document.getElementById("tabfmState").textContent = p.tabfm.status || "unknown";
      document.getElementById("rowCount").textContent = p.rows.length;
      document.getElementById("elapsedMs").textContent = `${Math.round(p.elapsed_ms)} ms`;
    }

    function renderDatasetControls() {
      const select = document.getElementById("datasetSelect");
      select.innerHTML = state.payload.datasets.map(ds => `<option value="${ds.id}">${ds.name}</option>`).join("");
      select.value = state.datasetId;
      select.onchange = event => {
        state.datasetId = event.target.value;
        const ds = datasetById(state.datasetId);
        state.sampleSize = ds.sample_sizes[0];
        renderAll();
      };
    }

    function renderDatasetStats() {
      const ds = datasetById(state.datasetId);
      document.getElementById("datasetRows").textContent = ds.rows;
      document.getElementById("datasetFeatures").textContent = ds.features;
      document.getElementById("datasetTask").textContent = ds.task === "classification" ? "Class" : "Reg";
      document.getElementById("lineTitle").textContent = `${ds.name}: ${ds.task === "classification" ? "Accuracy" : "R2"} vs Sample Size`;
      document.getElementById("metricBadge").textContent = ds.task === "classification" ? "accuracy" : "r2";
    }

    function renderSizeButtons() {
      const ds = datasetById(state.datasetId);
      const grid = document.getElementById("sizeGrid");
      grid.innerHTML = ds.sample_sizes.map(size => `
        <button class="size-btn ${Number(size) === Number(state.sampleSize) ? "active" : ""}" data-size="${size}">${size}</button>
      `).join("");
      grid.querySelectorAll("button").forEach(btn => {
        btn.addEventListener("click", () => {
          state.sampleSize = Number(btn.dataset.size);
          renderAll();
        });
      });
    }

    function renderLineChart() {
      const ds = datasetById(state.datasetId);
      const sizes = ds.sample_sizes;
      const datasets = state.payload.models.map(model => {
        const points = sizes.map(size => {
          const row = state.payload.rows.find(item =>
            item.dataset_id === state.datasetId &&
            item.model_name === model.model_name &&
            Number(item.sample_size) === Number(size)
          );
          return row && row.status === "ok" ? row.primary_score : null;
        });
        return {
          label: model.model_name,
          data: points,
          borderColor: COLORS[model.model_name],
          backgroundColor: COLORS[model.model_name],
          pointRadius: 4,
          borderWidth: model.model_name === "TabFM" ? 4 : 2,
          tension: 0.25,
          spanGaps: false
        };
      });
      const ctx = document.getElementById("lineChart");
      if (state.lineChart) state.lineChart.destroy();
      state.lineChart = new Chart(ctx, {
        type: "line",
        data: { labels: sizes, datasets },
        options: {
          responsive: true,
          maintainAspectRatio: false,
          interaction: { mode: "nearest", intersect: false },
          scales: {
            x: { grid: { color: "rgba(255,255,255,0.06)" }, ticks: { color: "#a6adbb" }, title: { display: true, text: "Training rows", color: "#a6adbb" } },
            y: { grid: { color: "rgba(255,255,255,0.06)" }, ticks: { color: "#a6adbb" }, suggestedMin: 0, suggestedMax: 1 }
          },
          plugins: {
            legend: { labels: { color: "#dce0ea", usePointStyle: true, boxWidth: 8 } },
            tooltip: { callbacks: { label: ctx => `${ctx.dataset.label}: ${formatScore(ctx.parsed.y)}` } }
          }
        }
      });
    }

    function renderLeaderboard() {
      const rows = rowsForSelection().slice().sort((a, b) => {
        const av = a.primary_score === null ? -Infinity : Number(a.primary_score);
        const bv = b.primary_score === null ? -Infinity : Number(b.primary_score);
        return bv - av;
      });
      const body = document.getElementById("leaderRows");
      body.innerHTML = rows.map((row, index) => `
        <tr>
          <td class="${index === 0 && row.status === "ok" ? "leader" : ""}">${row.status === "ok" ? index + 1 : "-"}</td>
          <td>${row.model_name}</td>
          <td class="score">${formatScore(row.primary_score, row.primary_metric)}</td>
          <td><span class="status ${statusClass(row.status)}">${row.status}</span></td>
        </tr>
      `).join("");
    }

    function renderWinChart() {
      const labels = state.payload.win_rates.map(row => row.model_name);
      const values = state.payload.win_rates.map(row => Math.round(row.win_rate * 1000) / 10);
      const colors = labels.map(label => COLORS[label]);
      const ctx = document.getElementById("winChart");
      if (state.winChart) state.winChart.destroy();
      state.winChart = new Chart(ctx, {
        type: "bar",
        data: { labels, datasets: [{ label: "Win rate", data: values, backgroundColor: colors, borderWidth: 0 }] },
        options: {
          responsive: true,
          maintainAspectRatio: false,
          scales: {
            x: { grid: { display: false }, ticks: { color: "#a6adbb" } },
            y: { grid: { color: "rgba(255,255,255,0.06)" }, ticks: { color: "#a6adbb", callback: value => `${value}%` }, suggestedMin: 0, suggestedMax: 100 }
          },
          plugins: { legend: { display: false } }
        }
      });
    }

    function renderModelCards() {
      const list = document.getElementById("modelList");
      list.innerHTML = state.payload.models.map(model => {
        const isActive = state.selectedModel === model.model_name;
        const score = model.average_primary_score === null ? "-" : Number(model.average_primary_score).toFixed(3);
        const status = model.model_name === "TabFM" ? model.runtime_status : "ok";
        return `
          <article class="model-card ${isActive ? "active" : ""}" data-model="${model.model_name}">
            <h3><span>${model.model_name}</span><span class="status ${statusClass(status)}">${status}</span></h3>
            <p>${model.type} | avg score ${score} | measured rows ${model.measured_rows}</p>
          </article>
        `;
      }).join("");
      list.querySelectorAll(".model-card").forEach(card => {
        card.addEventListener("click", () => {
          const model = card.dataset.model;
          state.selectedModel = state.selectedModel === model ? null : model;
          renderModelCards();
          renderDetails();
        });
      });
    }

    function renderDetails() {
      document.getElementById("selectedModelBadge").textContent = state.selectedModel || "All models";
      let rows = rowsForDataset().slice().sort((a, b) => Number(a.sample_size) - Number(b.sample_size) || a.model_name.localeCompare(b.model_name));
      if (state.selectedModel) rows = rows.filter(row => row.model_name === state.selectedModel);
      const body = document.getElementById("detailRows");
      body.innerHTML = rows.map(row => {
        const secondary = row.task === "classification" ? row.f1 : row.r2;
        const title = row.note ? ` title="${String(row.note).replaceAll('"', "&quot;")}"` : "";
        return `
          <tr${title}>
            <td>${row.sample_size}</td>
            <td>${row.model_name}</td>
            <td class="score">${formatScore(row.primary_score, row.primary_metric)}</td>
            <td>${formatScore(secondary)}</td>
            <td>${formatMs(row.train_time_ms)}</td>
            <td>${formatMs(row.inference_time_ms)}</td>
            <td><span class="status ${statusClass(row.status)}">${row.status}</span></td>
          </tr>
        `;
      }).join("");
    }

    loadPayload().catch(error => {
      const box = document.getElementById("errorBox");
      box.hidden = false;
      box.textContent = error.message;
    });
  </script>
</body>
</html>
"""


if __name__ == "__main__":
    app.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", "7860")),
        quiet=True,
    )