import os
import sys

import numpy as np
import pandas as pd
import torch
import gradio as gr
import plotly.graph_objects as go
from plotly.subplots import make_subplots

# Use all available CPU cores (HF free Spaces are CPU-only).
try:
    torch.set_num_threads(max(1, os.cpu_count() or 1))
except Exception:
    pass

# Make the vendored `model/` package importable.
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from model import Kronos, KronosTokenizer, KronosPredictor  # noqa: E402

# label -> (tokenizer_repo, model_repo, max_context)
MODELS = {
    "Kronos-mini (4.1M, ctx 2048)": ("NeoQuasar/Kronos-Tokenizer-2k", "NeoQuasar/Kronos-mini", 2048),
    "Kronos-small (24.7M, ctx 512)": ("NeoQuasar/Kronos-Tokenizer-base", "NeoQuasar/Kronos-small", 512),
    "Kronos-base (102.3M, ctx 512)": ("NeoQuasar/Kronos-Tokenizer-base", "NeoQuasar/Kronos-base", 512),
}
DEFAULT_MODEL = "Kronos-small (24.7M, ctx 512)"
PRICE_COLS = ["open", "high", "low", "close"]
TIME_CANDIDATES = ["timestamps", "timestamp", "date", "datetime", "time"]

_predictors = {}  # cache loaded models so we only download/instantiate once


def get_predictor(choice):
    if choice not in _predictors:
        tok_repo, mdl_repo, max_ctx = MODELS[choice]
        tokenizer = KronosTokenizer.from_pretrained(tok_repo)
        model = Kronos.from_pretrained(mdl_repo)
        _predictors[choice] = KronosPredictor(model, tokenizer, device="cpu", max_context=max_ctx)
    return _predictors[choice]


def _read_df(csv_file):
    if csv_file is None:
        path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "sample_data.csv")
    else:
        path = csv_file if isinstance(csv_file, str) else getattr(csv_file, "name", csv_file)
    df = pd.read_csv(path)
    df.columns = [str(c).strip().lower() for c in df.columns]
    return df


def _get_timestamps(df):
    for c in TIME_CANDIDATES:
        if c in df.columns:
            ts = pd.to_datetime(df[c], errors="coerce")
            if ts.notna().any():
                return ts.reset_index(drop=True)
    # No timestamp column -> synthesize a regular 5-minute index.
    return pd.Series(pd.date_range("2024-01-01", periods=len(df), freq="5min"))


def forecast(model_choice, lookback, pred_len, T, top_p, sample_count, csv_file):
    lookback, pred_len, sample_count = int(lookback), int(pred_len), int(sample_count)

    df = _read_df(csv_file)
    missing = [c for c in PRICE_COLS if c not in df.columns]
    if missing:
        raise gr.Error(f"CSV must contain columns {PRICE_COLS}. Missing: {missing}")
    if "volume" not in df.columns:
        df["volume"] = 0.0
    if "amount" not in df.columns:
        df["amount"] = 0.0

    _, _, max_ctx = MODELS[model_choice]
    lookback = min(lookback, max_ctx)
    if len(df) < lookback + 1:
        raise gr.Error(f"Need at least {lookback + 1} rows of history; the file has only {len(df)}.")

    df = df.reset_index(drop=True)
    ts_all = _get_timestamps(df)

    x_df = df.loc[: lookback - 1, ["open", "high", "low", "close", "volume", "amount"]]
    x_ts = ts_all.loc[: lookback - 1]

    have_future = len(df) >= lookback + pred_len
    if have_future:
        y_ts = ts_all.loc[lookback : lookback + pred_len - 1].reset_index(drop=True)
    else:
        deltas = ts_all.diff().dropna()
        step = deltas.median() if len(deltas) else pd.Timedelta(minutes=5)
        last = ts_all.iloc[lookback - 1]
        y_ts = pd.Series([last + step * (i + 1) for i in range(pred_len)])

    predictor = get_predictor(model_choice)
    pred_df = predictor.predict(
        df=x_df, x_timestamp=x_ts, y_timestamp=y_ts, pred_len=pred_len,
        T=float(T), top_p=float(top_p), sample_count=sample_count, verbose=False,
    ).reset_index(drop=True)

    # Build interactive chart: history + forecast (+ actual future when available).
    hist_x = list(range(lookback))
    fut_x = list(range(lookback, lookback + pred_len))
    fig = make_subplots(
        rows=2, cols=1, shared_xaxes=True, vertical_spacing=0.07,
        row_heights=[0.7, 0.3], subplot_titles=("Close price", "Volume"),
    )
    fig.add_trace(go.Scatter(x=hist_x, y=df.loc[: lookback - 1, "close"],
                             name="History", line=dict(color="#1f77b4")), row=1, col=1)
    fig.add_trace(go.Scatter(x=fut_x, y=pred_df["close"],
                             name="Kronos forecast", line=dict(color="#d62728", width=2)), row=1, col=1)
    if have_future:
        fig.add_trace(go.Scatter(x=fut_x, y=df.loc[lookback : lookback + pred_len - 1, "close"],
                                 name="Actual (future)", line=dict(color="#2ca02c", dash="dot")), row=1, col=1)
    fig.add_trace(go.Scatter(x=hist_x, y=df.loc[: lookback - 1, "volume"],
                             line=dict(color="#1f77b4"), showlegend=False), row=2, col=1)
    fig.add_trace(go.Scatter(x=fut_x, y=pred_df["volume"],
                             line=dict(color="#d62728"), showlegend=False), row=2, col=1)
    fig.add_vline(x=lookback - 1, line=dict(color="gray", dash="dash"))
    fig.update_layout(height=600, hovermode="x unified",
                      legend=dict(orientation="h", yanchor="bottom", y=1.04, xanchor="left", x=0),
                      margin=dict(l=50, r=20, t=60, b=40))

    out = pred_df.copy()
    out.insert(0, "timestamp", list(y_ts))
    return fig, out.round(4)


DESCRIPTION = """
# 📈 Kronos — Financial K-line Forecasting

Interactive demo of [**Kronos**](https://github.com/shiyu-coder/Kronos), the first open-source
foundation model for financial candlesticks (K-lines), trained on data from 45+ global exchanges.

Upload an OHLCV CSV (or use the bundled sample), choose a model size, and generate a probabilistic
forecast of future candles. Required columns: **open, high, low, close**; *volume*, *amount* and a
*timestamp* column are optional.

> ⚠️ **Not financial advice.** This is a research/demo tool. Forecasts are probabilistic samples from
> a model and must not be used as the sole basis for any trading or investment decision.
"""


with gr.Blocks(title="Kronos Forecast") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column(scale=1):
            model_choice = gr.Dropdown(list(MODELS.keys()), value=DEFAULT_MODEL, label="Model")
            csv_file = gr.File(label="OHLCV CSV (optional — uses sample if empty)", file_types=[".csv"])
            lookback = gr.Slider(64, 512, value=256, step=8, label="Lookback (history length)")
            pred_len = gr.Slider(10, 120, value=30, step=5, label="Forecast length (candles)")
            with gr.Accordion("Sampling parameters", open=False):
                T = gr.Slider(0.1, 2.0, value=1.0, step=0.1, label="Temperature (T)")
                top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p (nucleus)")
                sample_count = gr.Slider(1, 5, value=1, step=1, label="Sample count (averaged)")
            run_btn = gr.Button("Generate forecast", variant="primary")
            gr.Markdown("*CPU inference: a 30-step forecast takes roughly 10-30s.*")
        with gr.Column(scale=2):
            plot = gr.Plot(label="Forecast")
            table = gr.Dataframe(label="Forecast values", wrap=True)

    run_btn.click(
        forecast,
        inputs=[model_choice, lookback, pred_len, T, top_p, sample_count, csv_file],
        outputs=[plot, table],
    )
    gr.Markdown("Model weights: [NeoQuasar on Hugging Face](https://huggingface.co/NeoQuasar) - "
                "Code: [shiyu-coder/Kronos](https://github.com/shiyu-coder/Kronos) (MIT)")


if __name__ == "__main__":
    demo.queue(max_size=8).launch(
        theme=gr.themes.Soft(),
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
    )