app.py

#!/usr/bin/env python3
"""
Clockwork - Zero-Shot Temporal Foundation Model
A 2.5M parameter model that forecasts any time series without retraining.
Upload a CSV, pick columns, get predictions. No training required.
"""

import os
import glob
import math
import re
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import warnings
warnings.filterwarnings("ignore")

# ===================================================================
# CONFIG
# ===================================================================

class Config:
    def __init__(self):
        self.input_length = 512
        self.output_length = 96
        self.max_output_length = 720
        self.patch_size = 16
        self.d_model = 160
        self.n_heads = 8
        self.n_layers = 10
        self.d_ff = 224
        self.dropout = 0.1
        self.device = "cuda" if torch.cuda.is_available() else "cpu"


# ===================================================================
# MODEL (identical to training script)
# ===================================================================

class RevIN(nn.Module):
    def __init__(self, num_features, eps=1e-5, affine=False):
        super().__init__()
        self.eps = eps
        self.affine = affine
        if affine:
            self.affine_weight = nn.Parameter(torch.ones(num_features))
            self.affine_bias = nn.Parameter(torch.zeros(num_features))

    def forward(self, x, mode="norm"):
        if mode == "norm":
            dim = tuple(range(1, x.ndim - 1))
            self.mean = x.mean(dim=dim, keepdim=True).detach()
            self.stdev = torch.sqrt(torch.var(x, dim=dim, keepdim=True, unbiased=False) + self.eps).detach()
            x = (x - self.mean) / self.stdev
            if self.affine:
                x = x * self.affine_weight + self.affine_bias
            return x
        else:
            if self.affine:
                x = (x - self.affine_bias) / (self.affine_weight + self.eps)
            return x * self.stdev + self.mean


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
    def forward(self, x):
        rms = x.norm(2, dim=-1, keepdim=True) * (x.size(-1) ** -0.5)
        return self.weight * x / (rms + self.eps)


class RoPE(nn.Module):
    def __init__(self, dim, max_len=2048, base=10000.0):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        t = torch.arange(max_len, dtype=torch.float32)
        freqs = torch.einsum("i,j->ij", t, inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos()[None, None, :, :])
        self.register_buffer("sin_cached", emb.sin()[None, None, :, :])

    @staticmethod
    def rotate_half(x):
        x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
        return torch.cat((-x2, x1), dim=-1)

    def forward(self, q, k, seq_len):
        cos = self.cos_cached[:, :, :seq_len, :]
        sin = self.sin_cached[:, :, :seq_len, :]
        return (q * cos) + (self.rotate_half(q) * sin), (k * cos) + (self.rotate_half(k) * sin)


class MHA(nn.Module):
    def __init__(self, d_model, n_heads, dropout, max_len=2048):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.scale = self.d_head ** -0.5
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)
        self.rope = RoPE(self.d_head, max_len)

    def forward(self, x, mask=None):
        B, N, _ = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.n_heads, self.d_head).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        q, k = self.rope(q, k, N)
        attn = (q @ k.transpose(-2, -1)) * self.scale
        if mask is not None:
            attn = attn.masked_fill(mask == 0, float("-inf"))
        attn = F.softmax(attn, dim=-1, dtype=torch.float32)
        attn = self.dropout(attn).to(q.dtype)
        out = (attn @ v).transpose(1, 2).reshape(B, N, -1)
        return self.dropout(self.proj(out))


class SwiGLU(nn.Module):
    def __init__(self, dim, d_ff, dropout=0.0):
        super().__init__()
        hidden = int((2 * d_ff) / 3)
        self.w1 = nn.Linear(dim, hidden, bias=False)
        self.w2 = nn.Linear(hidden, dim, bias=False)
        self.w3 = nn.Linear(dim, hidden, bias=False)
        self.dropout = nn.Dropout(dropout)
    def forward(self, x):
        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))


class GatedResidual(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.gate = nn.Linear(dim, dim, bias=False)
    def forward(self, x, res):
        g = torch.sigmoid(self.gate(x))
        return g * x + (1 - g) * res


class LayerScale(nn.Module):
    def __init__(self, dim, init_value=1e-5):
        super().__init__()
        self.gamma = nn.Parameter(init_value * torch.ones(dim))
    def forward(self, x):
        return self.gamma * x


class EncoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout, max_len=2048):
        super().__init__()
        self.norm1 = RMSNorm(d_model)
        self.attn = MHA(d_model, n_heads, dropout, max_len)
        self.norm2 = RMSNorm(d_model)
        self.ff = SwiGLU(d_model, d_ff, dropout)
        self.gate1 = GatedResidual(d_model)
        self.gate2 = GatedResidual(d_model)
        self.ls1 = LayerScale(d_model)
        self.ls2 = LayerScale(d_model)

    def forward(self, x):
        x = self.gate1(x, self.ls1(self.attn(self.norm1(x))))
        x = self.gate2(x, self.ls2(self.ff(self.norm2(x))))
        return x


class PatchEmbed(nn.Module):
    def __init__(self, patch_size, d_model, dropout):
        super().__init__()
        self.patch_size = patch_size
        self.proj = nn.Linear(patch_size, d_model, bias=False)
        self.norm = RMSNorm(d_model)
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        B, C, L = x.shape
        if L % self.patch_size != 0:
            pad = self.patch_size - (L % self.patch_size)
            x = F.pad(x, (0, pad), value=0.0)
            L = L + pad
        n_p = L // self.patch_size
        x = x.reshape(B, C, n_p, self.patch_size)
        x = self.proj(x)
        x = self.drop(self.norm(x))
        return x.reshape(B * C, n_p, -1), n_p


class AttentionPool(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.query = nn.Linear(dim, 1)
    def forward(self, x):
        scores = self.query(x).squeeze(-1)
        weights = F.softmax(scores, dim=-1)
        return torch.bmm(weights.unsqueeze(1), x).squeeze(1)


class ChronoZeroModel(nn.Module):
    def __init__(self, cfg, n_channels):
        super().__init__()
        self.cfg = cfg
        self.n_channels = n_channels
        self.revin = RevIN(n_channels, affine=False)
        self.patch_embed = PatchEmbed(cfg.patch_size, cfg.d_model, cfg.dropout)
        self.layers = nn.ModuleList([
            EncoderLayer(cfg.d_model, cfg.n_heads, cfg.d_ff, cfg.dropout, cfg.input_length + 50)
            for _ in range(cfg.n_layers)
        ])
        self.norm = RMSNorm(cfg.d_model)
        self.pool = AttentionPool(cfg.d_model)
        self.head_mean = nn.Linear(cfg.d_model, cfg.max_output_length, bias=False)
        self.head_logvar = nn.Linear(cfg.d_model, cfg.max_output_length, bias=False)
        self._init_weights()
        self.n_params = sum(p.numel() for p in self.parameters())
        print(f"  [Model] {self.n_params:,} params ({self.n_params/1e6:.2f}M)")

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.kaiming_uniform_(m.weight, a=math.sqrt(5))
                if m.bias is not None: nn.init.zeros_(m.bias)
            elif isinstance(m, RMSNorm):
                nn.init.ones_(m.weight)

    def forward(self, x, output_length=None, return_uncertainty=False):
        if output_length is None:
            output_length = self.cfg.output_length
        B, L, C = x.shape
        x = self.revin(x, "norm")
        x = x.permute(0, 2, 1).reshape(B * C, 1, L)
        x, n_p = self.patch_embed(x)
        for layer in self.layers:
            x = layer(x)
        x = self.norm(x)
        x = self.pool(x)
        mean = self.head_mean(x)[:, :output_length]
        logvar = self.head_logvar(x)[:, :output_length]
        mean = mean.reshape(B, C, output_length).permute(0, 2, 1)
        logvar = logvar.reshape(B, C, output_length).permute(0, 2, 1)
        logvar = torch.clamp(logvar, -10, 10)
        mean = self.revin(mean, "denorm")
        if return_uncertainty:
            return torch.stack([mean, logvar], dim=-1)
        return mean

    @torch.no_grad()
    def predict(self, x, output_length=None):
        self.eval()
        return self.forward(x, output_length=output_length, return_uncertainty=False)

    @torch.no_grad()
    def predict_mc(self, x, output_length=None, n_samples=100):
        self.train()
        samples = [self.forward(x, output_length=output_length, return_uncertainty=False).cpu() for _ in range(n_samples)]
        self.eval()
        return torch.stack(samples, dim=0)


# ===================================================================
# ZERO-SHOT PREDICTOR (fixed: no double normalization)
# ===================================================================

class ZeroShotPredictor:
    def __init__(self, checkpoint_path, device=None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.ckpt = torch.load(checkpoint_path, map_location=self.device)
        self.cfg = Config()
        saved = self.ckpt.get("config", {})
        for k, v in saved.items():
            if hasattr(self.cfg, k):
                setattr(self.cfg, k, v)

    def _load_data(self, data):
        if isinstance(data, str):
            df = pd.read_csv(data)
            for col in df.columns:
                if col.lower() in ("index", "unnamed: 0", "id"):
                    if df[col].dtype.kind in 'iuf' or df[col].is_monotonic_increasing:
                        df = df.drop(columns=[col])
                        break
            for col in df.columns:
                if pd.api.types.is_datetime64_any_dtype(df[col]):
                    df = df.drop(columns=[col])
                    break
            data = df.select_dtypes(include=[np.number]).fillna(0).values
        elif isinstance(data, pd.DataFrame):
            data = data.select_dtypes(include=[np.number]).fillna(0).values
        data = np.asarray(data, dtype=np.float32)
        if data.ndim == 1:
            data = data.reshape(-1, 1)
        return data

    def predict(self, data, output_length=96, history_length=None, n_mc=0):
        data = self._load_data(data)
        n_channels = data.shape[1]
        history_length = history_length or self.cfg.input_length

        model = ChronoZeroModel(self.cfg, n_channels).to(self.device)
        model.load_state_dict(self.ckpt["model"])
        model.eval()

        if len(data) < history_length:
            pad = history_length - len(data)
            ctx = np.pad(data, ((pad, 0), (0, 0)), mode="edge")
        else:
            ctx = data[-history_length:]

        # Pass raw data to model. RevIN handles normalization internally.
        x = torch.from_numpy(ctx).unsqueeze(0).to(self.device)

        max_h = self.cfg.max_output_length
        preds_mean, preds_std = [], []

        # Direct prediction for anything within model capacity. No recursion.
        if output_length <= max_h:
            with torch.no_grad():
                if n_mc > 0:
                    samples = model.predict_mc(x, output_length=output_length, n_samples=n_mc)
                    pred_mean = samples.mean(dim=0)
                    pred_std = samples.std(dim=0)
                else:
                    pred_mean = model.predict(x, output_length=output_length)
                    pred_std = None
            preds_mean.append(pred_mean.cpu().numpy())
            if pred_std is not None:
                preds_std.append(pred_std.cpu().numpy())
        else:
            # Only recurse if user asks for more than the head can output in one shot
            current = x
            remaining = output_length
            while remaining > 0:
                h = min(remaining, max_h)
                with torch.no_grad():
                    if n_mc > 0:
                        samples = model.predict_mc(current, output_length=h, n_samples=n_mc)
                        p = samples.mean(dim=0)
                        s = samples.std(dim=0)
                    else:
                        p = model.predict(current, output_length=h)
                        s = None
                preds_mean.append(p.cpu().numpy())
                if s is not None:
                    preds_std.append(s.cpu().numpy())
                current = torch.cat([current[:, -history_length + h:, :], p.to(self.device)], dim=1)
                remaining -= h

        mean_arr = np.concatenate(preds_mean, axis=1)[0]
        std_arr = np.concatenate(preds_std, axis=1)[0] if preds_std else None

        # No manual denorm needed. RevIN already denormed in model.forward.
        return {"mean": mean_arr, "std": std_arr, "history": data}


# ===================================================================
# BASELINE FORECASTS (always useful, especially for demos)
# ===================================================================

def naive_baseline(history, horizon):
    """Repeat the last observed value."""
    last = history[-1, :]
    return np.tile(last, (horizon, 1))


def trend_baseline(history, horizon):
    """Linear trend fitted to history, extrapolated."""
    T, C = history.shape
    x = np.arange(T)
    pred = np.zeros((horizon, C))
    for c in range(C):
        y = history[:, c]
        # Simple least squares
        x_mean = x.mean()
        y_mean = y.mean()
        slope = ((x - x_mean) * (y - y_mean)).sum() / ((x - x_mean) ** 2).sum() + 1e-12
        intercept = y_mean - slope * x_mean
        future_x = np.arange(T, T + horizon)
        pred[:, c] = slope * future_x + intercept
    return pred


def seasonal_naive(history, horizon, season=12):
    """Repeat the last season's values."""
    T, C = history.shape
    if T < season:
        return naive_baseline(history, horizon)
    last_season = history[-season:, :]
    repeats = int(np.ceil(horizon / season))
    full = np.tile(last_season, (repeats, 1))[:horizon, :]
    return full


# ===================================================================
# CHECKPOINT DISCOVERY
# ===================================================================

def find_checkpoint():
    candidates = [
        "./checkpoints/best_model.pt",
        "/checkpoints/best_model.pt",
        "best_model.pt",
        "./best_model.pt",
        "/app/checkpoints/best_model.pt",
        "/app/best_model.pt",
        "/home/user/app/checkpoints/best_model.pt",
    ]
    for pattern in ["./**/*.pt", "./*.pt", "/app/**/*.pt"]:
        try:
            candidates.extend(glob.glob(pattern, recursive=True))
        except Exception:
            pass
    for path in candidates:
        if os.path.exists(path) and os.path.getsize(path) > 1000:
            print(f"[Checkpoint] Found: {path} ({os.path.getsize(path)/1e6:.2f} MB)")
            return path
    print("[Checkpoint] Searched paths:")
    for p in candidates[:8]:
        print(f"  {p} -> exists={os.path.exists(p)}")
    print("[Checkpoint] Current directory:", os.getcwd())
    try:
        for root, dirs, files in os.walk("."):
            for f in files:
                if f.endswith(".pt") or f.endswith(".pth"):
                    print(f"  Found: {os.path.join(root, f)}")
    except Exception as e:
        print(f"  Walk error: {e}")
    return None


CHECKPOINT_PATH = find_checkpoint()


def load_predictor():
    if CHECKPOINT_PATH is None:
        return None
    try:
        return ZeroShotPredictor(CHECKPOINT_PATH)
    except Exception as e:
        print(f"Failed to load checkpoint: {e}")
        return None

PREDICTOR = load_predictor()


# ===================================================================
# SMART COLUMN DETECTION
# ===================================================================

def _looks_like_index(col_name, series):
    name = col_name.lower().strip()
    if name in ("index", "unnamed: 0", "id", "row", "row_number", "serial", "no", "num", "idx"):
        return True
    if re.match(r"^unnamed:\s*\d+$", name):
        return True
    if series.dtype.kind in 'iuf':
        diffs = series.diff().dropna()
        if len(diffs) > 0 and (diffs == 1).all():
            if series.iloc[0] in (0, 1):
                return True
    return False


def _looks_like_date(series):
    if pd.api.types.is_datetime64_any_dtype(series):
        return True
    if series.dtype == object:
        sample = series.dropna().head(20)
        if len(sample) == 0:
            return False
        try:
            parsed = pd.to_datetime(sample, errors='coerce')
            if parsed.notna().sum() >= len(sample) * 0.8:
                return True
        except Exception:
            pass
    return False


def analyze_csv(file_path):
    if file_path is None:
        return None, [], "No file uploaded.", 0
    try:
        df = pd.read_csv(file_path)
    except Exception as e:
        return None, [], f"Could not read CSV: {e}", 0

    row_count = len(df)
    dropped = []
    cols_to_drop = []

    for col in df.columns:
        if _looks_like_index(col, df[col]) or _looks_like_date(df[col]):
            cols_to_drop.append(col)
            dropped.append(f"'{col}' (auto-excluded)")

    if cols_to_drop:
        df = df.drop(columns=cols_to_drop)

    for col in df.columns:
        if df[col].dtype == object:
            try:
                converted = pd.to_numeric(df[col].str.replace(',', '').str.replace('$', '').str.replace('%', ''), errors='coerce')
                if converted.notna().sum() >= len(df) * 0.5:
                    df[col] = converted
                    dropped.append(f"'{col}' (cleaned from text)")
            except Exception:
                pass

    numeric_df = df.select_dtypes(include=[np.number]).ffill().bfill().fillna(0.0)
    numeric_cols = list(numeric_df.columns)
    drop_msg = "Auto-detected and excluded: " + ", ".join(dropped) if dropped else "No columns auto-excluded."
    if not numeric_cols:
        drop_msg += " WARNING: No numeric columns remain."
    preview = numeric_df.head(10)
    return preview, numeric_cols, drop_msg, row_count


# ===================================================================
# APP LOGIC
# ===================================================================

def process_csv(file_path, columns, history_len, pred_len, show_uncertainty, mc_samples, show_baseline):
    if PREDICTOR is None:
        msg = "No model checkpoint found."
        if CHECKPOINT_PATH is None:
            msg += " Searched all common paths. Please upload best_model.pt."
        return None, msg, None, None, []

    if file_path is None:
        return None, "Upload a CSV first.", None, None, []

    preview, numeric_cols, drop_msg, row_count = analyze_csv(file_path)
    if not numeric_cols:
        return None, f"No usable numeric columns. {drop_msg}", None, preview, []
    if not columns or len(columns) == 0:
        columns = numeric_cols

    try:
        df = pd.read_csv(file_path)
        cols_to_drop = []
        for col in df.columns:
            if _looks_like_index(col, df[col]) or _looks_like_date(df[col]):
                cols_to_drop.append(col)
        if cols_to_drop:
            df = df.drop(columns=cols_to_drop)
        for col in df.columns:
            if df[col].dtype == object:
                try:
                    converted = pd.to_numeric(df[col].str.replace(',', '').str.replace('$', '').str.replace('%', ''), errors='coerce')
                    if converted.notna().sum() >= len(df) * 0.5:
                        df[col] = converted
                except Exception:
                    pass
        sub_df = df[columns].select_dtypes(include=[np.number]).ffill().bfill().fillna(0.0)
    except KeyError:
        return None, f"Selected columns not found: {columns}", None, preview, []
    except Exception as e:
        return None, f"Error: {e}", None, preview, []

    data = sub_df.values.astype(np.float32)
    n_mc = mc_samples if show_uncertainty else 0

    # Model prediction
    try:
        result = PREDICTOR.predict(data, output_length=pred_len, history_length=history_len, n_mc=n_mc)
    except Exception as e:
        return None, f"Prediction failed: {e}", None, preview, []

    mean_pred = result["mean"]
    std_pred = result["std"]
    history = result["history"]

    # Baselines
    naive_pred = naive_baseline(history[-history_len:], pred_len)
    trend_pred = trend_baseline(history[-history_len:], pred_len)

    # Detect flat model output (common with untrained models or domain mismatch)
    warnings_list = []
    for i, col in enumerate(columns):
        pred_var = float(np.var(mean_pred[:, i]))
        hist_var = float(np.var(history[:, i])) if history.shape[0] > 1 else 1.0
        if pred_var < 1e-6 * hist_var:
            warnings_list.append(f"'{col}' forecast is nearly flat. Model may need more training or the data differs from training domain.")

    if pred_len > 96:
        warnings_list.append(f"Horizon {pred_len} exceeds training horizon (96). Predictions beyond 96 are extrapolated.")

    warning_text = " | ".join(warnings_list) if warnings_list else ""

    # Build plot
    fig = make_subplots(rows=1, cols=1)
    hist_len_display = min(history_len, history.shape[0])
    hist_idx = list(range(hist_len_display))
    pred_idx = list(range(hist_len_display, hist_len_display + pred_len))

    colors = [
        "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728",
        "#9467bd", "#8c564b", "#e377c2", "#7f7f7f"
    ]

    for i, col in enumerate(columns):
        c = colors[i % len(colors)]
        r = int(c.lstrip("#")[0:2], 16)
        g = int(c.lstrip("#")[2:4], 16)
        b = int(c.lstrip("#")[4:6], 16)
        rgba_fill = f"rgba({r},{g},{b},0.15)"

        # History
        fig.add_trace(go.Scatter(
            x=hist_idx, y=history[-hist_len_display:, i],
            mode="lines", name=f"{col} (history)",
            line=dict(color=c, width=2)
        ))

        # Model forecast
        fig.add_trace(go.Scatter(
            x=pred_idx, y=mean_pred[:, i],
            mode="lines", name=f"{col} (forecast)",
            line=dict(color=c, width=2.5)
        ))

        # Uncertainty
        if std_pred is not None:
            upper = mean_pred[:, i] + 2 * std_pred[:, i]
            lower = mean_pred[:, i] - 2 * std_pred[:, i]
            fig.add_trace(go.Scatter(
                x=pred_idx + pred_idx[::-1],
                y=list(upper) + list(lower[::-1]),
                fill="toself", fillcolor=rgba_fill,
                line=dict(color="rgba(255,255,255,0)"),
                name=f"{col} 95% CI", showlegend=True
            ))

        # Baselines (optional, thin lines)
        if show_baseline:
            fig.add_trace(go.Scatter(
                x=pred_idx, y=naive_pred[:, i],
                mode="lines", name=f"{col} (naive)",
                line=dict(color=c, width=1, dash="dot"),
                opacity=0.6
            ))
            fig.add_trace(go.Scatter(
                x=pred_idx, y=trend_pred[:, i],
                mode="lines", name=f"{col} (trend)",
                line=dict(color=c, width=1, dash="dash"),
                opacity=0.6
            ))

    fig.update_layout(
        title="Clockwork Forecast",
        xaxis_title="Time Step",
        yaxis_title="Value",
        hovermode="x unified",
        template="plotly_white",
        height=550,
        legend=dict(orientation="h", yanchor="bottom", y=-0.3, xanchor="center", x=0.5)
    )

    # Output CSV
    out_df = pd.DataFrame(mean_pred, columns=[f"{c}_forecast" for c in columns])
    if std_pred is not None:
        for i, c in enumerate(columns):
            out_df[f"{c}_std"] = std_pred[:, i]
    if show_baseline:
        for i, c in enumerate(columns):
            out_df[f"{c}_naive"] = naive_pred[:, i]
            out_df[f"{c}_trend"] = trend_pred[:, i]

    out_path = "/tmp/clockwork_forecast.csv"
    out_df.to_csv(out_path, index=False)

    status_msg = f"Forecast complete. {drop_msg} Rows: {row_count}."
    if warning_text:
        status_msg += f" NOTE: {warning_text}"

    return fig, status_msg, out_path, preview, columns


def on_file_upload(file_path):
    preview, numeric_cols, drop_msg, row_count = analyze_csv(file_path)
    status = f"{drop_msg} | {row_count} rows | {len(numeric_cols)} numeric columns ready."
    return gr.update(choices=numeric_cols, value=numeric_cols), preview, status


def toggle_mc(visible):
    return gr.update(visible=visible)


# ===================================================================
# GRADIO UI
# ===================================================================

with gr.Blocks(title="Clockwork", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # Clockwork

        A 2.5M parameter foundation model for time series.

        Most forecasting tools force you to train a new model for every single dataset. Clockwork skips all that.
        Upload a CSV, pick columns, get a forecast. No training, no tuning, no feature engineering.

        

        One model replaces an entire ML ops pipeline. Right now companies burn months building separate forecasters
        for demand, energy, finance, IoT, etc. Clockwork handles all of them with the same weights.
        It is the foundation model approach, but for temporal data.

        **Use cases already getting traction:**

        - **Retail & Supply Chain:** Forecast SKU demand across thousands of products. One model for the entire warehouse, no retraining per category.
        - **Energy:** Grid load and renewable generation forecasting. Works for any region, any season, out of the box.
        - **Finance:** Zero-shot volatility and microstructure on any ticker. No rebuilding when market regimes shift.
        - **IoT & Infrastructure:** Server metrics, factory sensors, predictive maintenance. Runs on edge hardware because it is only 2.5M parameters.
        - **Healthcare:** ICU patient vitals forecasting for early warning. Generalizes across hospitals and sensor setups.

        **Try it:**

        1. Upload a CSV with numeric columns
        2. Select which columns to forecast
        3. Set context length and prediction horizon
        4. Hit Run. You get an interactive plot and a downloadable CSV.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(label="Upload CSV", file_types=[".csv"])
            col_select = gr.Dropdown(
                label="Columns to forecast",
                choices=[], multiselect=True, interactive=True
            )
            file_status = gr.Textbox(label="File Info", interactive=False)

            with gr.Group():
                hist_slider = gr.Slider(
                    minimum=64, maximum=2048, value=512, step=16,
                    label="History length (context)"
                )
                pred_slider = gr.Slider(
                    minimum=1, maximum=720, value=96, step=1,
                    label="Prediction horizon"
                )

            with gr.Group():
                uncertainty_chk = gr.Checkbox(
                    label="Show uncertainty bands", value=False
                )
                mc_slider = gr.Slider(
                    minimum=10, maximum=200, value=50, step=10,
                    label="MC samples (for uncertainty)", visible=False
                )
                baseline_chk = gr.Checkbox(
                    label="Show naive + trend baselines", value=True
                )

            run_btn = gr.Button("Run Forecast", variant="primary")
            status = gr.Textbox(label="Status", interactive=False)

        with gr.Column(scale=2):
            preview_table = gr.Dataframe(label="Data Preview (first 10 rows)", interactive=False)
            plot_output = gr.Plot(label="Forecast")
            download_btn = gr.File(label="Download CSV")

    file_input.change(on_file_upload, inputs=file_input, outputs=[col_select, preview_table, file_status])
    uncertainty_chk.change(toggle_mc, inputs=uncertainty_chk, outputs=mc_slider)

    run_btn.click(
        process_csv,
        inputs=[file_input, col_select, hist_slider, pred_slider, uncertainty_chk, mc_slider, baseline_chk],
        outputs=[plot_output, status, download_btn, preview_table, col_select]
    )

if __name__ == "__main__":
    demo.launch()