app.py

"""
Kronos-base forecasting app for Hugging Face Spaces (CPU or GPU).

IMPORTANT:
This file needs `model.py` (or the `model/` package) from the official
Kronos GitHub repo sitting in the SAME folder as this app.py:
    https://github.com/shiyu-coder/Kronos

NOTE on storage: prediction-tracking below is kept in memory for the
lifetime of this running Space (session-scoped). Hugging Face's FREE
Space disk is ephemeral — if the Space restarts/rebuilds, this history
resets. That's a platform limitation, not a bug in this app.
"""

import time
import tempfile

import gradio as gr
import numpy as np
import pandas as pd
import plotly.graph_objects as go
import torch
import yfinance as yf

from model import Kronos, KronosTokenizer, KronosPredictor

# ---------------------------------------------------------------------------
# Load model once at startup (auto-picks GPU if the Space has one, else CPU)
# ---------------------------------------------------------------------------
torch.set_num_threads(2)

DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
MAX_CONTEXT = 512

print(f"Device: {DEVICE}")
print("Loading tokenizer...")
tokenizer = KronosTokenizer.from_pretrained("NeoQuasar/Kronos-Tokenizer-base")

print("Loading Kronos-base model...")
model = Kronos.from_pretrained("NeoQuasar/Kronos-base")

predictor = KronosPredictor(model, tokenizer, device=DEVICE, max_context=MAX_CONTEXT)
print("Model ready.")


# ---------------------------------------------------------------------------
# Timezone: everything from Yahoo Finance is normalized to IST (UTC+5:30)
# ---------------------------------------------------------------------------
IST = "Asia/Kolkata"


def _now_ist():
    return pd.Timestamp.now(tz="UTC").tz_convert(IST).tz_localize(None)


def _to_ist_naive(ts_series):
    ts_series = pd.to_datetime(ts_series)
    if ts_series.dt.tz is not None:
        ts_series = ts_series.dt.tz_convert(IST)
    else:
        ts_series = ts_series.dt.tz_localize("UTC").dt.tz_convert(IST)
    return ts_series.dt.tz_localize(None)


# ---------------------------------------------------------------------------
# Symbol / timeframe catalogs (Yahoo Finance tickers)
# ---------------------------------------------------------------------------
CRYPTO_SYMBOLS = {
    "BTC/USD": "BTC-USD", "ETH/USD": "ETH-USD", "BNB/USD": "BNB-USD",
    "SOL/USD": "SOL-USD", "XRP/USD": "XRP-USD", "DOGE/USD": "DOGE-USD",
    "ADA/USD": "ADA-USD", "AVAX/USD": "AVAX-USD", "LTC/USD": "LTC-USD",
    "MATIC/USD": "MATIC-USD",
}

FOREX_SYMBOLS = {
    "EUR/USD": "EURUSD=X", "GBP/USD": "GBPUSD=X", "USD/JPY": "USDJPY=X",
    "USD/CHF": "USDCHF=X", "AUD/USD": "AUDUSD=X", "USD/CAD": "USDCAD=X",
    "NZD/USD": "NZDUSD=X", "EUR/JPY": "EURJPY=X", "GBP/JPY": "GBPJPY=X",
    "EUR/GBP": "EURGBP=X",
}

TIMEFRAMES = {
    "1 Minute": ("1m", "7d"), "5 Minutes": ("5m", "60d"),
    "15 Minutes": ("15m", "60d"), "30 Minutes": ("30m", "60d"),
    "1 Hour": ("60m", "730d"), "1 Day": ("1d", "5y"),
}

INTERVAL_TO_FREQ = {
    "1m": "1min", "5m": "5min", "15m": "15min",
    "30m": "30min", "60m": "1h", "1d": "1D",
}

INTERVAL_TO_TIMEDELTA = {
    "1m": pd.Timedelta(minutes=1), "5m": pd.Timedelta(minutes=5),
    "15m": pd.Timedelta(minutes=15), "30m": pd.Timedelta(minutes=30),
    "60m": pd.Timedelta(hours=1), "1d": pd.Timedelta(days=1),
}

_yf_cache = {}
_YF_CACHE_TTL = 20  # seconds


# ---------------------------------------------------------------------------
# Yahoo Finance fetch (light cache + retry)
# ---------------------------------------------------------------------------
def _fetch_yf_history(ticker, interval, period):
    cache_key = (ticker, interval, period)
    now = time.time()
    cached = _yf_cache.get(cache_key)
    if cached and (now - cached[0] < _YF_CACHE_TTL):
        return cached[1].copy()

    last_err = None
    for attempt in range(3):
        try:
            raw = yf.Ticker(ticker).history(period=period, interval=interval)
            if raw is not None and not raw.empty:
                _yf_cache[cache_key] = (now, raw)
                return raw.copy()
            last_err = ValueError("Khaali data mila")
        except Exception as e:
            last_err = e
        time.sleep(1.5 * (attempt + 1))

    raise RuntimeError(
        f"Yahoo Finance se '{ticker}' ka data nahi mil paaya ({last_err}). "
        f"Thodi der (30-60s) baad dubara try karo."
    )


def _load_context_from_yf(asset_type, symbol_label, timeframe_label, lookback, custom_ticker=""):
    if custom_ticker and custom_ticker.strip():
        ticker = custom_ticker.strip().upper()
    else:
        symbols = CRYPTO_SYMBOLS if asset_type == "Crypto" else FOREX_SYMBOLS
        ticker = symbols[symbol_label]

    interval, period = TIMEFRAMES[timeframe_label]
    raw = _fetch_yf_history(ticker, interval, period)

    raw = raw.reset_index()
    time_col = raw.columns[0]
    raw = raw.rename(columns={
        time_col: "timestamps", "Open": "open", "High": "high",
        "Low": "low", "Close": "close", "Volume": "volume",
    })
    raw["timestamps"] = _to_ist_naive(raw["timestamps"])  # <-- IST here
    keep = [c for c in ["timestamps", "open", "high", "low", "close", "volume"] if c in raw.columns]
    raw = raw[keep]

    df = _clean_ohlcv(raw, freq_hint=INTERVAL_TO_FREQ[interval])
    if len(df) < lookback:
        raise ValueError(
            f"{ticker} ({timeframe_label}) ke liye cleaning ke baad sirf {len(df)} valid "
            f"candles bache, {lookback} chahiye. Chhota lookback try karo ya bada timeframe."
        )
    df = df.tail(lookback).reset_index(drop=True)
    return df, INTERVAL_TO_FREQ[interval], ticker, interval, period


# ---------------------------------------------------------------------------
# Data cleaning
# ---------------------------------------------------------------------------
def _clean_ohlcv(df, freq_hint=None):
    df = df.copy()
    df = df.dropna(subset=["timestamps"])
    df = df.sort_values("timestamps")
    df = df.drop_duplicates(subset=["timestamps"], keep="last")

    freq = freq_hint or pd.infer_freq(df["timestamps"])
    if freq and len(df) > 1:
        full_index = pd.date_range(df["timestamps"].iloc[0], df["timestamps"].iloc[-1], freq=freq)
        df = df.set_index("timestamps").reindex(full_index)
        df.index.name = "timestamps"
        df[["open", "high", "low", "close"]] = df[["open", "high", "low", "close"]].ffill()
        if "volume" in df.columns:
            df["volume"] = df["volume"].fillna(0)
        df = df.reset_index()

    df = df.dropna(subset=["open", "high", "low", "close"]).reset_index(drop=True)
    if df.empty:
        return df

    df["high"] = df[["high", "open", "close"]].max(axis=1)
    df["low"] = df[["low", "open", "close"]].min(axis=1)

    pct_change = df["close"].pct_change().abs()
    df = df[(pct_change < 0.25) | pct_change.isna()].reset_index(drop=True)

    if "volume" in df.columns:
        df["volume"] = df["volume"].clip(lower=0)

    return df


def _postprocess_predictions(pred_df):
    pred_df = pred_df.copy()
    pred_df["high"] = pred_df[["high", "open", "close"]].max(axis=1)
    pred_df["low"] = pred_df[["low", "open", "close"]].min(axis=1)
    for c in ["volume", "amount"]:
        if c in pred_df.columns:
            pred_df[c] = pred_df[c].clip(lower=0)
    return pred_df


def _predict_with_confidence(x_df, feat_cols, x_timestamp, y_timestamp,
                              pred_len, temperature, top_p, sample_count, confidence_runs):
    runs = []
    for _ in range(max(1, int(confidence_runs))):
        pred = predictor.predict(
            df=x_df[feat_cols], x_timestamp=x_timestamp, y_timestamp=y_timestamp,
            pred_len=pred_len, T=float(temperature), top_p=float(top_p),
            sample_count=int(sample_count),
        )
        runs.append(_postprocess_predictions(pred))

    base = runs[0].copy()
    if len(runs) == 1:
        return base, None, None

    close_stack = pd.concat([r["close"].reset_index(drop=True) for r in runs], axis=1)
    mean_close = close_stack.mean(axis=1)
    std_close = close_stack.std(axis=1).fillna(0)

    base["close"] = mean_close.values
    lower = (mean_close - std_close).values
    upper = (mean_close + std_close).values

    rel_spread = std_close / (mean_close.abs() + 1e-9)
    confidence = (100 * (1 - rel_spread.clip(0, 1))).round(1)

    base["confidence_pct"] = confidence.values
    base["close_lower"] = lower
    base["close_upper"] = upper
    return base, lower, upper


def _build_plot(x_timestamp, x_close, pred_df, lower, upper, title_suffix=""):
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=x_timestamp, y=x_close, mode="lines",
                              name="History (close)", line=dict(color="#3b82f6")))

    if lower is not None and upper is not None:
        fig.add_trace(go.Scatter(x=pred_df.index, y=upper, mode="lines",
                                  line=dict(width=0), showlegend=False, hoverinfo="skip"))
        fig.add_trace(go.Scatter(x=pred_df.index, y=lower, mode="lines", line=dict(width=0),
                                  fill="tonexty", fillcolor="rgba(239,68,68,0.15)",
                                  name="Confidence band", hoverinfo="skip"))

    fig.add_trace(go.Scatter(x=pred_df.index, y=pred_df["close"], mode="lines",
                              name="Forecast (close)", line=dict(color="#ef4444")))

    fig.update_layout(
        title=f"Kronos-base Forecast {title_suffix} (IST)".strip(),
        xaxis_title="Time (IST)", yaxis_title="Price",
        hovermode="x unified", template="plotly_white",
        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
        margin=dict(l=40, r=20, t=60, b=40),
    )
    return fig


def _save_csv(df):
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv", prefix="kronos_forecast_")
    df.to_csv(tmp.name, index=False)
    return tmp.name


def _run_pipeline(x_df, freq, pred_len, temperature, top_p, sample_count, confidence_runs, title_suffix=""):
    x_timestamp = x_df["timestamps"]
    y_timestamp = pd.Series(pd.date_range(start=x_timestamp.iloc[-1], periods=pred_len + 1, freq=freq)[1:])
    feat_cols = [c for c in ["open", "high", "low", "close", "volume", "amount"] if c in x_df.columns]

    pred_df, lower, upper = _predict_with_confidence(
        x_df, feat_cols, x_timestamp, y_timestamp, pred_len,
        temperature, top_p, sample_count, confidence_runs,
    )

    fig = _build_plot(x_timestamp, x_df["close"], pred_df, lower, upper, title_suffix)

    table = pred_df.reset_index().rename(columns={"index": "timestamp"})
    front = ["timestamp", "open", "high", "low", "close"]
    if "close_lower" in table.columns:
        front += ["close_lower", "close_upper", "confidence_pct"]
    rest = [c for c in table.columns if c not in front]
    table = table[front + rest]

    csv_path = _save_csv(table)
    last_price = x_df["close"].iloc[-1]
    return fig, table, csv_path, last_price


# ---------------------------------------------------------------------------
# Live prediction tracking (session-scoped in-memory log) + verification
# ---------------------------------------------------------------------------
_prediction_log = []
_pred_id_counter = 0


def _log_prediction(ticker, interval, period, timeframe_label, anchor_close, pred_table):
    global _pred_id_counter
    _pred_id_counter += 1
    has_conf = "close_lower" in pred_table.columns
    rows = []
    for i, row in enumerate(pred_table.itertuples(index=False), start=1):
        rows.append({
            "step": i,
            "timestamp": getattr(row, "timestamp"),
            "predicted_close": float(getattr(row, "close")),
            "close_lower": float(getattr(row, "close_lower")) if has_conf else None,
            "close_upper": float(getattr(row, "close_upper")) if has_conf else None,
            "resolved": False,
            "actual_close": None,
            "error_pct": None,
            "direction_correct": None,
            "in_range": None,
        })
    _prediction_log.append({
        "id": _pred_id_counter, "ticker": ticker, "interval": interval, "period": period,
        "timeframe_label": timeframe_label, "made_at": _now_ist(),
        "anchor_close": float(anchor_close), "rows": rows,
    })


def verify_predictions():
    now = _now_ist()

    for entry in _prediction_log:
        interval, period = entry["interval"], entry["period"]
        candle_len = INTERVAL_TO_TIMEDELTA[interval]

        due_exists = any((not r["resolved"]) and (r["timestamp"] + candle_len <= now) for r in entry["rows"])
        if not due_exists:
            continue

        try:
            raw = _fetch_yf_history(entry["ticker"], interval, period)
        except Exception:
            continue

        raw = raw.reset_index()
        time_col = raw.columns[0]
        raw = raw[[time_col, "Close"]].rename(columns={time_col: "timestamps", "Close": "close"})
        raw["timestamps"] = _to_ist_naive(raw["timestamps"])
        raw = raw.set_index("timestamps")

        prev_actual = entry["anchor_close"]
        for r in sorted(entry["rows"], key=lambda x: x["step"]):
            if r["resolved"]:
                prev_actual = r["actual_close"]
                continue
            if r["timestamp"] + candle_len > now:
                break  # candle not fully closed yet — stop, keep chain order intact

            match = raw.index[raw.index == r["timestamp"]]
            if len(match) == 0:
                break  # data not published by Yahoo yet — try again next click

            actual_close = float(raw.loc[r["timestamp"], "close"])
            r["actual_close"] = actual_close
            r["error_pct"] = abs(r["predicted_close"] - actual_close) / actual_close * 100
            pred_dir = 1 if r["predicted_close"] >= prev_actual else -1
            act_dir = 1 if actual_close >= prev_actual else -1
            r["direction_correct"] = (pred_dir == act_dir)
            if r["close_lower"] is not None:
                r["in_range"] = r["close_lower"] <= actual_close <= r["close_upper"]
            r["resolved"] = True
            prev_actual = actual_close

    detail_rows = []
    pending = 0
    for entry in _prediction_log:
        for r in entry["rows"]:
            if r["resolved"]:
                detail_rows.append({
                    "made_at (IST)": entry["made_at"].strftime("%Y-%m-%d %H:%M"),
                    "ticker": entry["ticker"], "timeframe": entry["timeframe_label"],
                    "candle_#": r["step"], "timestamp (IST)": r["timestamp"].strftime("%Y-%m-%d %H:%M"),
                    "predicted_close": round(r["predicted_close"], 4),
                    "actual_close": round(r["actual_close"], 4),
                    "error_%": round(r["error_pct"], 3),
                    "direction": "✅ Correct" if r["direction_correct"] else "❌ Wrong",
                    "in_range": ("✅" if r["in_range"] else "❌") if r["in_range"] is not None else "-",
                })
            else:
                pending += 1

    if not detail_rows:
        return (f"Abhi verify karne layak koi candle nahi hai. Pending: {pending} "
                f"(ya time nahi aaya, ya Yahoo pe data abhi publish nahi hua)."), None, None

    detail_df = pd.DataFrame(detail_rows).sort_values(["made_at (IST)", "candle_#"], ascending=[False, True])

    summary_df = (
        pd.DataFrame(detail_rows)
        .groupby("candle_#")
        .agg(
            candles_checked=("candle_#", "count"),
            direction_accuracy_pct=("direction", lambda s: round(100 * (s == "✅ Correct").mean(), 1)),
            avg_error_pct=("error_%", "mean"),
        )
        .reset_index()
    )
    summary_df["avg_error_pct"] = summary_df["avg_error_pct"].round(3)

    overall_acc = round(100 * (detail_df["direction"] == "✅ Correct").mean(), 1)
    overall_err = round(detail_df["error_%"].mean(), 3)
    status = (
        f"✅ {len(detail_df)} candles verify hui | Overall direction accuracy: {overall_acc}% "
        f"| Avg error: {overall_err}% | Pending: {pending}"
    )
    return status, summary_df, detail_df


# ---------------------------------------------------------------------------
# Button handlers
# ---------------------------------------------------------------------------
def run_forecast_csv(csv_file, lookback, pred_len, temperature, top_p, sample_count, confidence_runs):
    try:
        if csv_file is None:
            return None, "CSV file upload karo pehle.", None, None

        lookback, pred_len = int(lookback), int(pred_len)
        if lookback > MAX_CONTEXT:
            lookback = MAX_CONTEXT

        df = pd.read_csv(csv_file.name)
        df.columns = [str(c).strip().lower() for c in df.columns]
        rename_map = {"timestamp": "timestamps", "time": "timestamps", "date": "timestamps",
                      "datetime": "timestamps", "open_time": "timestamps"}
        df = df.rename(columns={k: v for k, v in rename_map.items() if k in df.columns})

        required_cols = {"timestamps", "open", "high", "low", "close"}
        missing = required_cols - set(df.columns)
        if missing:
            return None, f"CSV mein ye columns missing hain: {missing}. Mile: {list(df.columns)}", None, None

        df["timestamps"] = pd.to_datetime(df["timestamps"])  # CSV timestamps used as-is (no tz assumed)
        freq = pd.infer_freq(df["timestamps"]) or "1min"
        df = _clean_ohlcv(df, freq_hint=freq)

        if len(df) < lookback:
            return None, f"Cleaning ke baad sirf {len(df)} valid rows hain, {lookback} chahiye.", None, None

        df = df.tail(lookback).reset_index(drop=True)
        fig, table, csv_path, last_price = _run_pipeline(
            df, freq, pred_len, temperature, top_p, sample_count, confidence_runs
        )
        return fig, f"Forecast ho gaya (CSV se). Last close: {last_price:.4f}", table, csv_path
    except Exception as e:
        return None, f"Error: {e}", None, None


def run_forecast_live(asset_type, symbol_label, timeframe_label, custom_ticker,
                       lookback, pred_len, temperature, top_p, sample_count, confidence_runs):
    try:
        lookback, pred_len = int(lookback), int(pred_len)
        if lookback > MAX_CONTEXT:
            lookback = MAX_CONTEXT

        df, freq, ticker, interval, period = _load_context_from_yf(
            asset_type, symbol_label, timeframe_label, lookback, custom_ticker
        )
        fig, table, csv_path, last_price = _run_pipeline(
            df, freq, pred_len, temperature, top_p, sample_count, confidence_runs,
            title_suffix=f"— {ticker} LIVE",
        )
        last_time = df["timestamps"].iloc[-1]

        _log_prediction(ticker, interval, period, timeframe_label, last_price, table)

        status = f"🔴 LIVE: {ticker} | Last price: {last_price:.5f} (as of {last_time.strftime('%Y-%m-%d %H:%M:%S')} IST)"
        live_price_text = f"{ticker}:  {last_price:.5f}   |   {last_time.strftime('%Y-%m-%d %H:%M:%S')} IST"
        return fig, status, table, csv_path, live_price_text
    except Exception as e:
        return None, f"Error: {e}", None, None, None


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="Kronos-base Forecaster") as demo:
    gr.Markdown(
        """
        # 📈 Kronos-base Forecaster — Prediction Only
        Future OHLC candle forecasting. **No trading signals, no buy/sell advice,
        no alerts, no portfolio tracking.** Saare live times **IST (UTC+5:30)** mein hain.
        """
    )

    data_source = gr.Radio(
        ["📁 Upload CSV", "🌐 Yahoo Finance (Live)"],
        value="🌐 Yahoo Finance (Live)", label="Data Source",
    )

    with gr.Row():
        with gr.Column(scale=1):
            with gr.Group(visible=False) as csv_group:
                csv_input = gr.File(label="OHLCV CSV Upload karo", file_types=[".csv"])
                csv_btn = gr.Button("📊 CSV Se Forecast Karo", variant="primary")

            with gr.Group(visible=True) as yf_group:
                asset_type = gr.Radio(["Crypto", "Forex"], value="Crypto", label="Asset Type")
                symbol_dd = gr.Dropdown(choices=list(CRYPTO_SYMBOLS.keys()), value="BTC/USD", label="Symbol")
                timeframe_dd = gr.Dropdown(choices=list(TIMEFRAMES.keys()), value="5 Minutes", label="Timeframe")
                custom_ticker = gr.Textbox(
                    label="Ya Custom Yahoo Finance Ticker (optional)",
                    placeholder="jaise: SOL-USD, GBPJPY=X",
                )
                live_btn = gr.Button("🔴 Live Prediction Nikalo", variant="primary")
                live_price_output = gr.Textbox(label="💰 Live Price (IST)", interactive=False)
                gr.Markdown("_Yahoo free data hai — rate-limit ho to 30-60s baad dubara try karo._")

            lookback_input = gr.Slider(64, 512, value=400, step=8, label="Lookback (max 512)")
            pred_len_input = gr.Slider(10, 240, value=60, step=10, label="Prediction Length")
            temperature_input = gr.Slider(0.1, 1.5, value=1.0, step=0.1, label="Temperature (T)")
            top_p_input = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
            sample_count_input = gr.Slider(1, 5, value=1, step=1, label="Sample Count (per-run averaging)")
            confidence_runs_input = gr.Slider(
                1, 5, value=1, step=1,
                label="Confidence Runs (N) — >1 se Expected Range + Confidence % milega (slow hoga)",
            )

        with gr.Column(scale=2):
            plot_output = gr.Plot(label="Forecast Chart")
            status_output = gr.Textbox(label="Status", interactive=False)
            table_output = gr.Dataframe(label="Forecast Data")
            csv_output = gr.File(label="⬇️ Forecast CSV Download")

    gr.Markdown("---\n## ✅ Live Prediction Verification")
    gr.Markdown(
        "Jab bhi **Live Prediction Nikalo** click karte ho, wo prediction yahan log ho jaati hai. "
        "Jitne candles ka time (IST mein) ab tak pura ho chuka hai, unhe neeche button se **automatically** "
        "actual Yahoo Finance data se verify kiya ja sakta hai — candle-by-candle, kaunsi sahi kaunsi galat."
    )
    verify_btn = gr.Button("✅ Verify Predictions Now", variant="secondary")
    verify_status = gr.Textbox(label="Verification Status", interactive=False)
    with gr.Row():
        verify_summary = gr.Dataframe(label="Candle-Position Wise Accuracy (1st candle, 2nd, ... last)")
        verify_detail = gr.Dataframe(label="Har Candle Ka Detail (Predicted vs Actual)")

    def toggle_source(source):
        is_csv = source.startswith("📁")
        return gr.update(visible=is_csv), gr.update(visible=not is_csv)

    data_source.change(toggle_source, inputs=[data_source], outputs=[csv_group, yf_group])

    def update_symbols(asset):
        choices = list(CRYPTO_SYMBOLS.keys()) if asset == "Crypto" else list(FOREX_SYMBOLS.keys())
        return gr.update(choices=choices, value=choices[0])

    asset_type.change(update_symbols, inputs=[asset_type], outputs=[symbol_dd])

    csv_btn.click(
        fn=run_forecast_csv,
        inputs=[csv_input, lookback_input, pred_len_input, temperature_input, top_p_input,
                sample_count_input, confidence_runs_input],
        outputs=[plot_output, status_output, table_output, csv_output],
    )

    live_btn.click(
        fn=run_forecast_live,
        inputs=[asset_type, symbol_dd, timeframe_dd, custom_ticker,
                lookback_input, pred_len_input, temperature_input, top_p_input,
                sample_count_input, confidence_runs_input],
        outputs=[plot_output, status_output, table_output, csv_output, live_price_output],
    )

    verify_btn.click(fn=verify_predictions, inputs=[], outputs=[verify_status, verify_summary, verify_detail])

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