"""
Play with google/tabfm-1.0.0-pytorch (TabFM): a zero-shot tabular foundation
model. Upload your own train/test CSVs, or try a built-in example, and get
predictions with no training, no hyperparameter search.
"""
import json
import os
import time

import gradio as gr
import spaces
import numpy as np
import pandas as pd
import torch
from huggingface_hub import snapshot_download, hf_hub_download
from safetensors.torch import load_file

from tabfm import TabFMClassifier, TabFMRegressor
from tabfm.src.pytorch.tabfm_v1_0_0 import TabFM_HF, HF_REPO_ID

_MODEL_CACHE = {}
_STOCK_DATA_CACHE = {}

LOOKBACK = 20

DATASETS = {
    "original": {
        "label": "Original — ticker holdout (TSLA / NFLX / AMD / WMT)",
        "repo": "ewinregirgojr/minicpm5-stock-v2-forward-return",
        "train_file": "train_minicpm5_v2.jsonl",
        "test_file": "val_minicpm5_v2.jsonl",
        "feature_keys": ["rsi", "vol_ratio", "volat20", "mom5", "mom20", "rel20"],
        "test_tickers": ["AMD", "NFLX", "TSLA", "WMT"],
    },
    "no_rel20": {
        "label": "No rel20 — last 30 per ticker (all 20 tickers)",
        "repo": "saviopreis/tabfm-stock-no-rel20",
        "train_file": "train.jsonl",
        "test_file": "test.jsonl",
        "feature_keys": ["rsi", "vol_ratio", "volat20", "mom5", "mom20"],
        "test_tickers": ["AAPL", "AMD", "AMZN", "BAC", "CSCO", "DIS", "GOOGL", "HD",
                         "JPM", "MA", "META", "MSFT", "NFLX", "NVDA", "PG",
                         "TSLA", "UNH", "V", "WMT", "XOM"],
    },
    "time_split": {
        "label": "All features — last 30 per ticker (all 20 tickers)",
        "repo": "saviopreis/tabfm-stock-time-split",
        "train_file": "train.jsonl",
        "test_file": "test.jsonl",
        "feature_keys": ["rsi", "vol_ratio", "volat20", "mom5", "mom20", "rel20"],
        "test_tickers": ["AAPL", "AMD", "AMZN", "BAC", "CSCO", "DIS", "GOOGL", "HD",
                         "JPM", "MA", "META", "MSFT", "NFLX", "NVDA", "PG",
                         "TSLA", "UNH", "V", "WMT", "XOM"],
    },
}


def _load_tabfm(model_type: str):
    """Load TabFM weights directly from safetensors, bypassing the pip load()
    which hardcodes pytorch_model.bin (not present in the HF repo)."""
    if model_type in _MODEL_CACHE:
        return _MODEL_CACHE[model_type]

    base_path = snapshot_download(repo_id=HF_REPO_ID)
    sub_path = os.path.join(base_path, model_type)

    with open(os.path.join(sub_path, "config.json")) as f:
        cfg = json.load(f)

    # Normalise config: HF Hub uses task="classification", local uses is_classifier
    if "task" in cfg:
        cfg["is_classifier"] = cfg.pop("task") == "classification"
    for key in ("model_type", "version", "framework", "frameworks"):
        cfg.pop(key, None)

    model = TabFM_HF(**cfg)
    state_dict = load_file(os.path.join(sub_path, "model.safetensors"), device="cpu")
    model.load_state_dict(state_dict, strict=True)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = model.to(torch.bfloat16).to(device)
    model.eval()

    _MODEL_CACHE[model_type] = model
    return model


@spaces.GPU(duration=120)
def _run(train_df, test_df, target_col, task, n_estimators, progress=gr.Progress()):
    if target_col not in train_df.columns:
        raise gr.Error(f"Target column '{target_col}' not found in training data.")

    Xtr = train_df.drop(columns=[target_col]).to_numpy()
    ytr = train_df[target_col].to_numpy()

    has_target_in_test = target_col in test_df.columns
    Xte = test_df.drop(columns=[target_col]).to_numpy() if has_target_in_test else test_df.to_numpy()

    progress(0.1, desc="Loading TabFM weights (cached after first run)...")
    t0 = time.time()
    model = _load_tabfm(task)
    load_s = time.time() - t0

    progress(0.4, desc="Running zero-shot in-context inference...")
    t0 = time.time()
    if task == "classification":
        clf = TabFMClassifier(model=model, random_state=42, n_estimators=int(n_estimators))
        clf.fit(Xtr, ytr)
        preds = clf.predict(Xte)
        proba = clf.predict_proba(Xte)
        classes = list(clf.classes_)
        pred_idx = [classes.index(p) for p in preds]
        prob_pred = np.array([proba[i, pred_idx[i]] for i in range(len(preds))])
    else:
        reg = TabFMRegressor(model=model, random_state=42, n_estimators=int(n_estimators))
        reg.fit(Xtr, ytr)
        preds = reg.predict(Xte)
        prob_pred = None
    infer_s = time.time() - t0

    out = test_df.copy()
    out["tabfm_prediction"] = preds
    if prob_pred is not None:
        out["prob_%"] = (prob_pred * 100).round(1)

    summary = f"Loaded weights in {load_s:.1f}s (cached after first run). Inference took {infer_s:.1f}s for {len(Xte)} test rows.\n"
    if has_target_in_test:
        y_true = test_df[target_col].to_numpy()
        if task == "classification":
            acc = (preds == y_true).mean()
            summary += f"Accuracy on your labeled test rows: {acc*100:.2f}%"
        else:
            mae = np.mean(np.abs(preds - y_true))
            summary += f"Mean absolute error on your labeled test rows: {mae:.4f}"
    return out, summary


def _rows_to_feature_matrix(rows, feature_keys):
    X, y, meta = [], [], []
    for row in rows:
        f = row["features"]
        X.append([f[k] for k in feature_keys] + f["rets20"])
        y.append(row["label"])
        meta.append({"ticker": row["ticker"], "date": row["date"], "fwd_return": row.get("fwd_return")})
    return np.array(X), np.array(y), meta


def _load_stock_dataset(dataset_key="original"):
    if dataset_key in _STOCK_DATA_CACHE:
        return _STOCK_DATA_CACHE[dataset_key]
    cfg = DATASETS[dataset_key]
    train_path = hf_hub_download(cfg["repo"], cfg["train_file"], repo_type="dataset")
    test_path  = hf_hub_download(cfg["repo"], cfg["test_file"],  repo_type="dataset")
    train_rows = [json.loads(l) for l in open(train_path, encoding="utf-8")]
    test_rows  = [json.loads(l) for l in open(test_path,  encoding="utf-8")]
    Xtr, ytr, _ = _rows_to_feature_matrix(train_rows, cfg["feature_keys"])
    data = {"Xtr": Xtr, "ytr": ytr, "test_rows": test_rows, "feature_keys": cfg["feature_keys"]}
    _STOCK_DATA_CACHE[dataset_key] = data
    return data


def _calculate_rsi(prices, window=14):
    deltas = np.diff(prices)
    seed = deltas[:window + 1]
    up = seed[seed >= 0].sum() / window
    down = -seed[seed < 0].sum() / window
    rs = up / down if down != 0 else 0
    rsi = np.zeros_like(prices, dtype=float)
    rsi[:window] = 100. - 100. / (1. + rs)
    for i in range(window, len(prices)):
        delta = deltas[i - 1]
        upval = delta if delta > 0 else 0.
        downval = -delta if delta < 0 else 0.
        up = (up * (window - 1) + upval) / window
        down = (down * (window - 1) + downval) / window
        rs = up / down if down != 0 else 0
        rsi[i] = 100. - 100. / (1. + rs)
    return rsi


@spaces.GPU(duration=300)
def _run_stock_backtest(dataset_key, ticker, n_estimators, progress=gr.Progress()):
    progress(0.1, desc=f"Loading dataset '{dataset_key}'...")
    data = _load_stock_dataset(dataset_key)
    test_rows = [r for r in data["test_rows"] if r["ticker"] == ticker]
    if not test_rows:
        raise gr.Error(f"No test rows for {ticker} in dataset '{dataset_key}'.")

    Xva, yva_str, meta = _rows_to_feature_matrix(test_rows, data["feature_keys"])
    ytr_bin = (data["ytr"] == "BUY").astype(int)

    progress(0.3, desc="Loading TabFM weights...")
    model = _load_tabfm("classification")
    clf = TabFMClassifier(model=model, random_state=42, n_estimators=int(n_estimators))

    progress(0.6, desc=f"Predicting {len(Xva)} test rows for {ticker}...")
    clf.fit(data["Xtr"], ytr_bin)
    proba   = clf.predict_proba(Xva)
    pred_bin = (proba[:, 1] > proba[:, 0]).astype(int)
    pred    = np.where(pred_bin == 1, "BUY", "SELL")
    prob_pred = np.where(pred_bin == 1, proba[:, 1], proba[:, 0])

    acc = (pred == yva_str).mean()
    out = pd.DataFrame({
        "date":                    [m["date"] for m in meta],
        "actual_label":            yva_str,
        "tabfm_prediction":        pred,
        "prob_%":                  (prob_pred * 100).round(1),
        "actual_5d_fwd_return_%":  [round(m["fwd_return"] * 100, 2) for m in meta],
        "correct":                 pred == yva_str,
    }).sort_values("date")

    ds_label = DATASETS[dataset_key]["label"]
    summary = (
        f"Dataset: {ds_label}\n"
        f"Ticker: {ticker} | Test rows: {len(Xva)} | Features: {data['feature_keys']}\n"
        f"Accuracy: {acc*100:.2f}%"
    )
    return out, summary


def _build_live_features(ticker):
    import yfinance as yf

    hist = yf.Ticker(ticker).history(period="1y", interval="1d", auto_adjust=True)
    spy = yf.Ticker("SPY").history(period="1y", interval="1d", auto_adjust=True)
    if hist.empty or len(hist) < 30:
        raise gr.Error(f"Couldn't fetch enough price history for '{ticker}' from yfinance.")

    hist.index = hist.index.tz_localize(None).normalize()
    spy.index = spy.index.tz_localize(None).normalize()
    df = hist.join(spy["Close"].rename("spy_close"), how="inner")

    c = df["Close"].values
    v = df["Volume"].values
    spyc = df["spy_close"].values
    if len(c) < 25:
        raise gr.Error(f"Not enough overlapping SPY/{ticker} history to build features.")

    ret1 = np.zeros(len(c))
    ret1[1:] = c[1:] / c[:-1] - 1.0
    rsi = _calculate_rsi(c)
    vol_ma = pd.Series(v).rolling(20).mean().values
    vol_ratio = v / vol_ma
    volat20 = pd.Series(ret1).rolling(20).std().values
    mom5 = c[-1] / c[-6] - 1.0
    mom20 = c[-1] / c[-21] - 1.0
    spy20 = spyc[-1] / spyc[-21] - 1.0
    rel20 = mom20 - spy20

    i = len(c) - 1
    feat_row = [rsi[i], vol_ratio[i], volat20[i], mom5, mom20, rel20] + list(ret1[i - LOOKBACK + 1:i + 1])
    as_of_date = df.index[-1].strftime("%Y-%m-%d")
    last_close = float(c[-1])
    return np.array(feat_row), as_of_date, last_close


LIVE_CONTEXT_SIZE = 600


def _live_context(data, seed=42):
    """Stratified subsample of the training context, used only for the live
    tab. TabFM's in-context inference cost scales with context size, and the
    full 5k-row train set takes several minutes for a single prediction --
    the published ~64% backtest number (Stock prediction backtest tab) still
    uses the full training set, this subsample is only for live-tab latency.
    """
    Xtr, ytr = data["Xtr"], data["ytr"]
    rng = np.random.default_rng(seed)
    idx_buy = np.where(ytr == "BUY")[0]
    idx_sell = np.where(ytr != "BUY")[0]
    n_each = LIVE_CONTEXT_SIZE // 2
    pick = np.concatenate([
        rng.choice(idx_buy, size=min(n_each, len(idx_buy)), replace=False),
        rng.choice(idx_sell, size=min(n_each, len(idx_sell)), replace=False),
    ])
    rng.shuffle(pick)
    return Xtr[pick], ytr[pick]


@spaces.GPU(duration=180)
def _run_live_prediction(ticker, n_estimators, progress=gr.Progress()):
    ticker = ticker.strip().upper()
    if not ticker:
        raise gr.Error("Enter a ticker symbol, e.g. NVDA.")

    progress(0.1, desc=f"Fetching recent {ticker} + SPY data from yfinance...")
    feat_row, as_of_date, last_close = _build_live_features(ticker)

    progress(0.3, desc="Loading training context + TabFM weights (first run per session is slowest)...")
    data = _load_stock_dataset("original")
    Xtr, ytr = _live_context(data)
    ytr_bin = (ytr == "BUY").astype(int)
    model = _load_tabfm("classification")
    clf = TabFMClassifier(model=model, random_state=42, n_estimators=int(n_estimators))

    progress(0.7, desc="Running TabFM zero-shot inference (this step alone can take 30s-2min on free CPU)...")
    clf.fit(Xtr, ytr_bin)
    proba = clf.predict_proba(feat_row.reshape(1, -1))[0]
    pred = "BUY" if proba[1] > proba[0] else "SELL"

    summary = (
        f"**{ticker}** as of {as_of_date} (last close ${last_close:.2f})\n\n"
        f"TabFM zero-shot prediction for the next 5 trading days: **{pred}**\n"
        f"(BUY probability: {proba[1]*100:.1f}% / SELL probability: {proba[0]*100:.1f}%)\n\n"
        f"⚠️ This is a live, unverified forecast of the *future* — there is no ground truth "
        f"to check it against yet. Our held-out backtest on quarantined tickers "
        f"(see the 'Stock prediction backtest' tab) scored ~64% accuracy, well above the "
        f"~53% ceiling classical models (GBM/RF/LogReg) found on the same features — "
        f"but past backtest accuracy is not a guarantee for any single live prediction. "
        f"Not financial advice."
    )
    return summary


def _load_example(name):
    from sklearn.datasets import load_breast_cancer, load_wine, load_diabetes

    if name == "Breast Cancer (classification)":
        d = load_breast_cancer(as_frame=True)
        df = d.frame
        task = "classification"
    elif name == "Wine (classification)":
        d = load_wine(as_frame=True)
        df = d.frame
        task = "classification"
    else:
        d = load_diabetes(as_frame=True)
        df = d.frame
        task = "regression"

    df = df.sample(frac=1.0, random_state=0).reset_index(drop=True)
    n_train = int(len(df) * 0.8)
    train_df, test_df = df.iloc[:n_train], df.iloc[n_train:]
    return train_df, test_df, "target" if "target" in df.columns else df.columns[-1], task


HEAD_HTML = """
<meta name="description" content="TabFM Playground: try Google's zero-shot tabular foundation model (TabFM), a TabPFN alternative, on your own data or a live stock-prediction demo. No training required.">
<meta name="author" content="saviopreis">
<meta property="article:modified_time" content="2026-07-06">
<script>document.documentElement.lang = "en";</script>
<script type="application/ld+json">
{
  "@context": "https://schema.org",
  "@type": "SoftwareApplication",
  "name": "TabFM Playground",
  "applicationCategory": "Machine Learning Tool",
  "operatingSystem": "Web",
  "description": "Zero-shot tabular foundation model (TabFM) demo from Google Research. Run classification/regression with no training, or try a live stock-prediction example.",
  "author": {"@type": "Person", "name": "saviopreis"},
  "dateModified": "2026-07-06",
  "offers": {"@type": "Offer", "price": "0", "priceCurrency": "USD"}
}
</script>
<script type="application/ld+json">
{
  "@context": "https://schema.org",
  "@type": "FAQPage",
  "mainEntity": [
    {
      "@type": "Question",
      "name": "Is TabFM better than XGBoost?",
      "acceptedAnswer": {
        "@type": "Answer",
        "text": "It depends on the dataset. In this Space's stock-prediction demo, TabFM's zero-shot predictions beat GradientBoosting/RandomForest/LogisticRegression run on identical features (~64% vs ~53% held-out accuracy), but this is one dataset, not a universal claim."
      }
    },
    {
      "@type": "Question",
      "name": "Can I use TabFM commercially?",
      "acceptedAnswer": {
        "@type": "Answer",
        "text": "No, not with these weights. TabFM ships under the TabFM Non-Commercial License v1.0: testing, evaluation, and research only. A commercial license would need to come from Google directly."
      }
    },
    {
      "@type": "Question",
      "name": "Does this Space train a new model?",
      "acceptedAnswer": {
        "@type": "Answer",
        "text": "No. TabFM never trains - it reads your training rows as context and predicts test rows in a single forward pass, the same in-context mechanism popularized by TabPFN."
      }
    },
    {
      "@type": "Question",
      "name": "Why is the first prediction slow?",
      "acceptedAnswer": {
        "@type": "Answer",
        "text": "The first call in a session loads an approximately 11GB checkpoint into memory. Predictions can take from 10 seconds to a couple of minutes on free CPU hardware depending on data size and ensemble size."
      }
    }
  ]
}
</script>
"""

with gr.Blocks(title="TabFM Playground") as demo:
    gr.Markdown(
        """
# TabFM Playground — Google's tabular foundation model, zero-shot

**TabFM is a tabular foundation model from Google Research** — a TabPFN-style
zero-shot alternative that predicts your target column via in-context learning
instead of training a new model. Give it [google/tabfm-1.0.0-pytorch](https://huggingface.co/google/tabfm-1.0.0-pytorch)
a training table and a test table and it makes predictions in a single forward
pass: no training loop, no hyperparameter search.

⚠️ Runs on free CPU hardware here, so inference can take anywhere from
~10s to a couple of minutes depending on data size and ensemble size.
Duplicate this Space with a GPU for much faster results.
        """
    )

    with gr.Tab("Quick demo"):
        example_dd = gr.Dropdown(
            ["Breast Cancer (classification)", "Wine (classification)", "Diabetes (regression)"],
            value="Breast Cancer (classification)",
            label="Built-in example dataset",
        )
        n_est_demo = gr.Slider(1, 8, value=2, step=1, label="Ensemble size (n_estimators, lower = faster)")
        run_demo_btn = gr.Button("Run TabFM on this example", variant="primary")
        demo_out = gr.Dataframe(label="Predictions")
        demo_summary = gr.Textbox(label="Summary", interactive=False)

        def _run_demo(name, n_est):
            train_df, test_df, target_col, task = _load_example(name)
            return _run(train_df, test_df, target_col, task, n_est)

        run_demo_btn.click(_run_demo, inputs=[example_dd, n_est_demo], outputs=[demo_out, demo_summary])

    with gr.Tab("Bring your own CSV"):
        gr.Markdown(
            "Upload a training CSV (with your target column) and a test CSV "
            "(target column optional — include it to see accuracy/MAE)."
        )
        train_file = gr.File(label="Training CSV", file_types=[".csv"])
        test_file = gr.File(label="Test CSV", file_types=[".csv"])
        target_col_tb = gr.Textbox(label="Target column name")
        task_dd = gr.Radio(["classification", "regression"], value="classification", label="Task type")
        n_est_custom = gr.Slider(1, 8, value=2, step=1, label="Ensemble size (n_estimators, lower = faster)")
        run_custom_btn = gr.Button("Run TabFM", variant="primary")
        custom_out = gr.Dataframe(label="Predictions")
        custom_summary = gr.Textbox(label="Summary", interactive=False)

        def _run_custom(train_f, test_f, target_col, task, n_est):
            if train_f is None or test_f is None:
                raise gr.Error("Please upload both a training CSV and a test CSV.")
            train_df = pd.read_csv(train_f.name)
            test_df = pd.read_csv(test_f.name)
            return _run(train_df, test_df, target_col, task, n_est)

        run_custom_btn.click(
            _run_custom,
            inputs=[train_file, test_file, target_col_tb, task_dd, n_est_custom],
            outputs=[custom_out, custom_summary],
        )

    with gr.Tab("Stock prediction backtest"):
        gr.Markdown(
            """
BUY/SELL prediction on 5-day forward returns using causal features (RSI, momentum,
volatility, last-20-day returns). Classical models topped out at ~51-53%.
Compare three dataset configurations to evaluate how robust TabFM really is:

| Dataset | Split | rel20 |
|---|---|---|
| **Original** | 4 tickers fully quarantined from training | ✅ |
| **No rel20** | Last 30 samples/ticker in test | ❌ removed |
| **All features** | Last 30 samples/ticker in test | ✅ |
            """
        )
        dataset_dd = gr.Dropdown(
            choices=[
                ("Original — ticker holdout (TSLA / NFLX / AMD / WMT)", "original"),
                ("No rel20 — last 30 per ticker (all 20 tickers)",       "no_rel20"),
                ("All features — last 30 per ticker (all 20 tickers)",   "time_split"),
            ],
            value="original",
            label="Dataset",
        )
        ticker_dd = gr.Dropdown(
            choices=DATASETS["original"]["test_tickers"],
            value=DATASETS["original"]["test_tickers"][0],
            label="Ticker",
        )
        n_est_stock = gr.Slider(1, 8, value=4, step=1, label="Ensemble size (n_estimators, lower = faster)")
        run_stock_btn = gr.Button("Run backtest", variant="primary")
        stock_out = gr.Dataframe(label="Predictions vs actual outcome")
        stock_summary = gr.Textbox(label="Summary", interactive=False)

        def _update_tickers(dataset_key):
            tickers = DATASETS[dataset_key]["test_tickers"]
            return gr.update(choices=tickers, value=tickers[0])

        dataset_dd.change(_update_tickers, inputs=[dataset_dd], outputs=[ticker_dd])
        run_stock_btn.click(
            _run_stock_backtest,
            inputs=[dataset_dd, ticker_dd, n_est_stock],
            outputs=[stock_out, stock_summary],
        )

    with gr.Tab("Live stock prediction (yfinance)"):
        gr.Markdown(
            """
Pulls **real-time recent price data via yfinance**, builds the exact same causal
features as the backtest above (nothing computed with future data), and asks TabFM
for a live BUY/SELL call — using the same v2 training set as in-context data, on a
ticker it may never have seen at all.

⚠️ Research demo only. This is a live forecast with no ground truth yet — **not
financial advice**, and past backtest accuracy does not guarantee any single
prediction.

🕒 **Expect 30 seconds to a few minutes on free CPU hardware** — the progress
bar above the button will update as it fetches data, loads weights, and runs
inference. It has not frozen; TabFM's zero-shot inference cost scales with
how much context data it's given.
            """
        )
        live_ticker_tb = gr.Textbox(label="Ticker symbol", placeholder="e.g. NVDA, AAPL, SPY")
        n_est_live = gr.Slider(1, 8, value=2, step=1, label="Ensemble size (n_estimators, lower = faster)")
        run_live_btn = gr.Button("Get live prediction", variant="primary")
        live_out = gr.Markdown()

        run_live_btn.click(_run_live_prediction, inputs=[live_ticker_tb, n_est_live], outputs=[live_out])

    gr.Markdown(
        """
---
Model: [google/tabfm-1.0.0-pytorch](https://huggingface.co/google/tabfm-1.0.0-pytorch) ·
Code: [google-research/tabfm](https://github.com/google-research/tabfm) ·
License: **TabFM Non-Commercial License v1.0** — the model weights are for non-commercial testing, evaluation, and research only. This Space is a non-commercial demo; commercial/production use requires a separate license from Google.
        """
    )

if __name__ == "__main__":
    demo.queue().launch(ssr_mode=False, head=HEAD_HTML)