"""
Kronos Monte-Carlo 預測 —— Gradio + ZeroGPU 版。

原本是 Docker + Flask 跑在 HF 免費 CPU 上，一次 6-15 分鐘。
ZeroGPU 只支援 Gradio SDK，所以外殼改寫成 Gradio；預測的數學一行都沒動。

三個 ZeroGPU 的硬性要求，對應到下面的寫法：
  1. Gradio SDK          → README.md 的 sdk: gradio，本檔用 gr.Interface
  2. 模型 module level 載入 → 下面的 tokenizer / model / predictor 都在模組層
  3. 推論包 @spaces.GPU   → _monte_carlo()

上游的 Kronos 套件（github.com/shiyu-coder/Kronos）沒有 setup.py 也沒有
pyproject.toml，所以 pip 裝不了。原本的 Dockerfile 是 build 時 git clone 進去，
Gradio Space 沒有 Dockerfile，改成啟動時把 model/ 那三個檔抓下來。
"""
import json
import os
import sys
import time
import urllib.request
from datetime import timedelta

# ─── 先把上游的 Kronos 套件放進 sys.path ─────────────────────────────────────
# 上游 repo 沒有 setup.py / pyproject.toml，pip 裝不了，所以直接抓 model/ 那三個檔
# （__init__.py / kronos.py / module.py，合計約 54KB）。
#
# ⚠️ shiyu-coder/Kronos 的預設分支是 **master**，不是 main。填錯會在開機時 404。
# KRONOS_REF 沒設的話會依序試 master、main；正式用建議直接釘成 commit SHA，
# 免得上游一改動你的 Space 就跟著壞。
KRONOS_FILES = ("__init__.py", "kronos.py", "module.py")
KRONOS_DIR = os.getenv("KRONOS_DIR", "/tmp/kronos_src")
_ref_env = os.getenv("KRONOS_REF", "").strip()
KRONOS_REFS = [_ref_env] if _ref_env else ["master", "main"]

def _raw_url(ref: str, fname: str) -> str:
    return f"https://raw.githubusercontent.com/shiyu-coder/Kronos/{ref}/model/{fname}"

def _bootstrap_kronos() -> None:
    pkg = os.path.join(KRONOS_DIR, "model")
    os.makedirs(pkg, exist_ok=True)

    missing = [f for f in KRONOS_FILES
               if not (os.path.exists(os.path.join(pkg, f))
                       and os.path.getsize(os.path.join(pkg, f)) > 100)]
    if not missing:
        if KRONOS_DIR not in sys.path:
            sys.path.insert(0, KRONOS_DIR)
        return

    errors = []
    for ref in KRONOS_REFS:
        try:
            for fname in missing:
                url = _raw_url(ref, fname)
                print(f"[bootstrap] fetching {fname} @ {ref}")
                with urllib.request.urlopen(url, timeout=30) as resp:
                    body = resp.read()
                # raw.githubusercontent 對不存在的檔案是回 404，但保險起見也擋一下
                # 「拿到 HTML 錯誤頁卻當成 Python」這種更難查的情況。
                if len(body) < 100 or body.lstrip()[:1] == b"<":
                    raise RuntimeError(f"{url} 回傳的內容不像 Python（{len(body)} bytes）")
                with open(os.path.join(pkg, fname), "wb") as fh:
                    fh.write(body)
            print(f"[bootstrap] ok @ {ref}")
            if KRONOS_DIR not in sys.path:
                sys.path.insert(0, KRONOS_DIR)
            return
        except Exception as exc:  # noqa: BLE001 —— 換下一個 ref 再試
            print(f"[bootstrap] ref={ref} 失敗：{exc}")
            errors.append(f"{ref}: {exc}")

    raise RuntimeError(
        "抓不到 Kronos 原始碼。試過的 ref：" + "、".join(KRONOS_REFS) + "\n"
        + "\n".join(errors)
        + "\n請確認 https://github.com/shiyu-coder/Kronos 的分支名稱，"
          "或用環境變數 KRONOS_REF 指定分支／commit SHA。"
    )

_bootstrap_kronos()

# spaces 要在 torch 之前 import（ZeroGPU 會去接管 CUDA 的初始化）
import spaces  # noqa: E402
import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402
import torch  # noqa: E402
import gradio as gr  # noqa: E402

from model import Kronos, KronosTokenizer, KronosPredictor  # noqa: E402

# ─── 設定 ────────────────────────────────────────────────────────────────────
# 模型尺寸做成環境變數，不寫死。三種選擇（參數量 / context / 搭配的 tokenizer）：
#   NeoQuasar/Kronos-mini    4.1M   2048   NeoQuasar/Kronos-Tokenizer-2k
#   NeoQuasar/Kronos-small  24.7M    512   NeoQuasar/Kronos-Tokenizer-base
#   NeoQuasar/Kronos-base  102.3M    512   NeoQuasar/Kronos-Tokenizer-base   ← 預設
# 換 mini 的話 tokenizer 一定要跟著換，只改 MODEL_REPO 會對不起來。
MODEL_REPO = os.getenv("KRONOS_MODEL", "NeoQuasar/Kronos-base")
TOKENIZER_REPO = os.getenv("KRONOS_TOKENIZER", "NeoQuasar/Kronos-Tokenizer-base")
MAX_CONTEXT = int(os.getenv("KRONOS_MAX_CONTEXT", "512"))

# ZeroGPU 上，module level 的時候 GPU 還沒配給你，torch.cuda.is_available() 會回
# False —— 用它判斷會誤選 cpu，GPU 就白開了。所以看 HF 注入的旗標，不看 torch。
ZERO_GPU = os.getenv("SPACES_ZERO_GPU", "").lower() in ("1", "true")
if ZERO_GPU:
    DEVICE = "cuda"
else:
    DEVICE = os.getenv("KRONOS_DEVICE") or ("cuda" if torch.cuda.is_available() else "cpu")

# @spaces.GPU 的額度是按**宣告值**整包扣的（先扣不退），不是按實際跑多久。
# 免費帳號一天 300 秒：宣告 60 秒 → 一天 5 次；宣告 30 秒 → 10 次。
# 所以這個值要盡量貼近實測時間，不要留一大堆保險。
GPU_DURATION = int(os.getenv("KRONOS_GPU_DURATION", "60"))

print(f"[init] model={MODEL_REPO} tokenizer={TOKENIZER_REPO} "
      f"device={DEVICE} zero_gpu={ZERO_GPU} max_context={MAX_CONTEXT}")

tokenizer = KronosTokenizer.from_pretrained(TOKENIZER_REPO)
model = Kronos.from_pretrained(MODEL_REPO)
predictor = KronosPredictor(
    model=model, tokenizer=tokenizer, device=DEVICE, max_context=MAX_CONTEXT
)
print("[init] predictor ready")


@spaces.GPU(duration=GPU_DURATION)
def _monte_carlo(df, ts, y_timestamp, pred_len, sample_count, T, top_p):
    """實際吃 GPU 的那一段。跑 sample_count 條各自獨立的軌跡。

    這裡刻意維持「跑 N 次、每次 sample_count=1」的寫法 —— 扇形圖需要 N 條
    各自獨立的路徑，predict() 的 sample_count 參數是內部合成後只回一條。
    """
    t0 = time.time()
    paths_close = []
    paths_ohlc = []
    for _ in range(sample_count):
        pred_df = predictor.predict(
            df=df[["open", "high", "low", "close", "volume", "amount"]],
            x_timestamp=pd.Series(ts),
            y_timestamp=y_timestamp,
            pred_len=pred_len,
            T=T,
            top_p=top_p,
            sample_count=1,
            verbose=False,
        )
        paths_close.append(pred_df["close"].astype(float).tolist())
        paths_ohlc.append([
            {
                "ts": int(y_timestamp.iloc[i].timestamp() * 1000),
                "open":  float(pred_df["open"].iloc[i]),
                "high":  float(pred_df["high"].iloc[i]),
                "low":   float(pred_df["low"].iloc[i]),
                "close": float(pred_df["close"].iloc[i]),
            }
            for i in range(pred_len)
        ])
    return paths_close, paths_ohlc, time.time() - t0


def predict(k_lines_json, pred_len=24, sample_count=10, interval_ms=3600000,
            T=1.0, top_p=0.9):
    """Gradio 的對外函式。回傳的 JSON 欄位跟舊版 Flask 完全一致。"""
    try:
        k_lines = json.loads(k_lines_json) if isinstance(k_lines_json, str) else k_lines_json
    except (TypeError, ValueError) as exc:
        return {"error": f"k_lines 不是合法的 JSON：{exc}"}

    if not k_lines or len(k_lines) < 50:
        return {"error": "need at least 50 k_lines"}

    pred_len = int(pred_len)
    sample_count = int(sample_count)
    interval_ms = int(interval_ms)
    T = float(T)
    top_p = float(top_p)

    # k_lines: [[ts_ms, open, high, low, close, volume], ...]
    arr = np.array(k_lines, dtype=float)
    ts = pd.to_datetime(arr[:, 0], unit="ms")
    df = pd.DataFrame({
        "open":   arr[:, 1],
        "high":   arr[:, 2],
        "low":    arr[:, 3],
        "close":  arr[:, 4],
        "volume": arr[:, 5] if arr.shape[1] > 5 else np.zeros(len(arr)),
        "amount": np.zeros(len(arr)),  # some versions expect this col
    })
    last_close = float(df["close"].iloc[-1])

    # Baseline volatility: realized std over a window of the SAME length as
    # the forecast horizon (pred_len). Comparing a 360-bar std to a 24-bar
    # forecast std is unfair and systematically makes vol_amp_prob ~= 0.
    tail_closes = df["close"].values[-(pred_len + 1):]
    hist_returns = np.diff(np.log(tail_closes))
    hist_std = float(np.std(hist_returns))
    # Guard — if history was perfectly flat, fall back to full-series std
    if hist_std == 0.0:
        full_rets = np.diff(np.log(df["close"].values))
        hist_std = float(np.std(full_rets)) if len(full_rets) else 1e-9

    last_ts = ts[-1]
    y_timestamp = pd.Series(
        [last_ts + timedelta(milliseconds=interval_ms * (i + 1)) for i in range(pred_len)]
    )

    paths_close, paths_ohlc, elapsed = _monte_carlo(
        df, ts, y_timestamp, pred_len, sample_count, T, top_p
    )

    finals = np.array([p[-1] for p in paths_close])
    upside_prob = float(np.mean(finals > last_close))

    path_stds = np.array([float(np.std(np.diff(np.log(np.array(p))))) for p in paths_close])
    # Ratio form: mean predicted vol / baseline vol. >1 means amplified, <1 dampened.
    vol_amp_ratio = float(np.mean(path_stds) / hist_std) if hist_std > 0 else 0.0
    # Probability form: fraction of paths more volatile than baseline (kept for reference).
    vol_amp_prob = float(np.mean(path_stds > hist_std))

    return {
        "ok": True,
        "model": MODEL_REPO,
        "device": DEVICE,
        "sample_count": sample_count,
        "pred_len": pred_len,
        "elapsed_s": round(elapsed, 2),
        "last_close": last_close,
        "hist_std": hist_std,
        "upside_prob": upside_prob,
        "vol_amp": vol_amp_ratio,        # legacy alias
        "vol_amp_ratio": vol_amp_ratio,
        "vol_amp_prob": vol_amp_prob,
        "mean_path_close": np.mean(np.array(paths_close), axis=0).tolist(),
        "paths": paths_ohlc,
        "future_ts_ms": [int(t.timestamp() * 1000) for t in y_timestamp],
    }


demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Textbox(label="k_lines (JSON)", lines=3,
                   info="[[ts_ms, open, high, low, close, volume], ...]，至少 50 根"),
        gr.Number(label="pred_len", value=24, precision=0),
        gr.Number(label="sample_count", value=10, precision=0),
        gr.Number(label="interval_ms", value=3600000, precision=0),
        gr.Number(label="T", value=1.0),
        gr.Number(label="top_p", value=0.9),
    ],
    outputs=gr.JSON(label="prediction"),
    api_name="predict",
    title="Kronos Monte-Carlo",
    description=(f"model={MODEL_REPO} · device={DEVICE} · "
                 f"gpu_duration={GPU_DURATION}s\n\n"
                 "API: POST /gradio_api/call/predict → 取 event_id → "
                 "GET /gradio_api/call/predict/{event_id}（SSE）"),
)

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)