app.py

import datetime as dt
import pandas as pd
import torch
import gradio as gr
import requests
import matplotlib
matplotlib.use("Agg")  # HF Space 같이 GUI 없는 환경에서 안전하게
import matplotlib.pyplot as plt

from chronos import BaseChronosPipeline  # pip install chronos-forecasting


# =============================
#   Chronos 모델 캐시/로더
# =============================
_PIPELINE_CACHE = {}


def get_pipeline(model_id: str, device: str = "cpu"):
    key = (model_id, device)
    if key not in _PIPELINE_CACHE:
        _PIPELINE_CACHE[key] = BaseChronosPipeline.from_pretrained(
            model_id,
            device_map=device,  # "cpu" / "cuda"
            torch_dtype=torch.float32 if device == "cpu" else torch.bfloat16,
        )
    return _PIPELINE_CACHE[key]


# =============================
#   Binance 전용 로더
# =============================
_BINANCE_INTERVAL = {"1d": "1d", "1h": "1h", "30m": "30m", "15m": "15m", "5m": "5m"}


def _yf_to_binance_symbol(ticker: str) -> str | None:
    """
    BTC-USD -> BTCUSDT, ETH-USD -> ETHUSDT ...
    그 외 형식은 None (현재는 -USD 코인만 지원)
    """
    t = ticker.upper().strip()
    if t.endswith("-USD") and len(t) >= 6:
        base = t[:-4]  # remove "-USD"
        return f"{base}USDT"
    return None


def _fetch_binance_klines(
    ticker: str, interval: str, start: str | None, end: str | None
) -> pd.Series:
    """
    Binance Klines (무인증)
    https://api.binance.com/api/v3/klines

    반환: pandas.Series(index=datetime, values=float close)
    """
    if interval not in _BINANCE_INTERVAL:
        raise ValueError("Binance는 해당 interval을 지원하지 않습니다.")

    symbol = _yf_to_binance_symbol(ticker)
    if not symbol:
        raise ValueError(
            "이 티커는 Binance 심볼로 변환할 수 없습니다. 예: BTC-USD, ETH-USD 형태만 지원합니다."
        )

    base = "https://api.binance.com/api/v3/klines"

    def to_ms(s: str) -> int:
        return int(pd.to_datetime(s).timestamp() * 1000)

    start_ms = to_ms(start) if start else None
    end_ms = to_ms(end) if end else None

    rows = []
    cur_start = start_ms

    while True:
        params = {
            "symbol": symbol,
            "interval": _BINANCE_INTERVAL[interval],
            "limit": 1000,
        }
        if cur_start is not None:
            params["startTime"] = cur_start
        if end_ms is not None:
            params["endTime"] = end_ms

        r = requests.get(base, params=params, timeout=30)
        r.raise_for_status()
        data = r.json()
        if not data:
            break

        rows.extend(data)

        # data[i] = [
        #   0 openTime,
        #   1 open,
        #   2 high,
        #   3 low,
        #   4 close,
        #   5 volume,
        #   6 closeTime,
        #   ...
        # ]
        last_close_time = data[-1][6]  # closeTime (ms)
        next_start = last_close_time + 1
        if cur_start is not None and next_start <= cur_start:
            break
        cur_start = next_start

        # 더 이상 1000개 안 나오면 마지막 페이지로 판단
        if len(data) < 1000:
            break

    if not rows:
        raise ValueError("Binance에서 가져온 데이터가 없습니다.")

    df = pd.DataFrame(
        rows,
        columns=[
            "openTime",
            "open",
            "high",
            "low",
            "close",
            "volume",
            "closeTime",
            "quoteAssetVolume",
            "numTrades",
            "takerBuyBase",
            "takerBuyQuote",
            "ignore",
        ],
    )
    df["ts"] = pd.to_datetime(df["closeTime"], unit="ms")
    s = df.set_index("ts")["close"].astype(float).sort_index()

    if start:
        s = s[s.index >= pd.to_datetime(start)]
    if end:
        s = s[s.index <= pd.to_datetime(end)]
    if s.empty:
        raise ValueError("Binance 시리즈가 비어 있습니다. 기간/간격을 다시 설정해 주세요.")

    return s


def load_close_series(
    ticker: str, start: str | None, end: str | None, interval: str = "1d"
) -> pd.Series:
    """
    Binance 전용 종가 시리즈 로더.
    입력: 티커 (예: BTC-USD, ETH-USD), 시작일, 종료일, 간격(1d/1h/30m/15m/5m)
    반환: pandas.Series (index=datetime, values=float close)
    """
    ticker = ticker.strip().upper()

    # 기본 기간 설정
    _start = start or "2017-01-01"
    _end = end or dt.date.today().isoformat()
    try:
        sdt = pd.to_datetime(_start)
        edt = pd.to_datetime(_end)
        if edt < sdt:
            sdt, edt = edt, sdt
        _start, _end = sdt.date().isoformat(), edt.date().isoformat()
    except Exception:
        # 잘못된 날짜 포맷이 들어오면 일단 기본값 사용
        _start, _end = "2017-01-01", dt.date.today().isoformat()

    # 오로지 Binance만 사용
    return _fetch_binance_klines(ticker, interval, _start, _end)


# =============================
#   예측 함수 (Gradio 핸들러)
# =============================
def run_forecast(ticker, start_date, end_date, horizon, model_id, device, interval):
    # 1) 시계열 로딩
    try:
        series = load_close_series(ticker, start_date, end_date, interval)
    except Exception as e:
        return None, pd.DataFrame(), f"데이터 로딩 오류 (Binance): {e}"

    # 2) 파이프라인 로딩
    try:
        pipe = get_pipeline(model_id, device)
    except Exception as e:
        return None, pd.DataFrame(), f"모델 로딩 오류: {e}"

    # 3) 예측 길이
    try:
        H = int(horizon)
        if H <= 0:
            raise ValueError("예측 스텝 H는 1 이상이어야 합니다.")
    except Exception:
        return None, pd.DataFrame(), "예측 스텝 H가 올바르지 않습니다."

    # 4) Chronos 입력
    context = torch.tensor(series.values, dtype=torch.float32)

    # 5) 예측: (num_series=1, num_quantiles=3, H) with q=[0.1, 0.5, 0.9]
    try:
        preds = pipe.predict(context=context, prediction_length=H)[0]
    except Exception as e:
        return None, pd.DataFrame(), f"예측 실행 오류: {e}"

    q10, q50, q90 = preds[0], preds[1], preds[2]

    # 6) 결과 표 (DataFrame)
    df_fcst = pd.DataFrame(
        {"q10": q10.numpy(), "q50": q50.numpy(), "q90": q90.numpy()},
        index=pd.RangeIndex(1, H + 1, name="step"),
    )

    # 7) 미래 x축 만들기
    freq_map = {"1d": "D", "1h": "H", "30m": "30T", "15m": "15T", "5m": "5T"}
    freq = freq_map.get(interval, "D")
    future_index = pd.date_range(series.index[-1], periods=H + 1, freq=freq)[1:]

    # 8) 그래프 그리기
    fig = plt.figure(figsize=(10, 4))
    plt.plot(series.index, series.values, label="history")
    plt.plot(future_index, q50.numpy(), label="forecast(q50)")
    plt.fill_between(
        future_index,
        q10.numpy(),
        q90.numpy(),
        alpha=0.2,
        label="q10–q90",
    )
    plt.title(f"{ticker} forecast by Chronos-Bolt (Binance, {interval}, H={H})")
    plt.legend()
    plt.tight_layout()

    note = "※ 데모 목적입니다. 투자 판단과 결과 책임은 전적으로 본인에게 있습니다."
    return fig, df_fcst, note


# =============================
#   Gradio UI
# =============================
with gr.Blocks(title="Chronos Crypto Forecast (Binance)") as demo:
    gr.Markdown("# Chronos 크립토 예측 데모 (Binance 전용)")
    gr.Markdown(
        "티커는 `BTC-USD`, `ETH-USD` 처럼 입력하면 내부에서 `BTCUSDT`, `ETHUSDT`로 변환해서 Binance에서 가격을 가져옵니다."
    )

    with gr.Row():
        ticker = gr.Textbox(
            value="BTC-USD",
            label="티커 (예: BTC-USD, ETH-USD)",
        )
        horizon = gr.Slider(
            5,
            365,
            value=90,
            step=1,
            label="예측 스텝 H (간격 단위와 동일)",
        )

    with gr.Row():
        start = gr.Textbox(
            value="2017-01-01",
            label="시작일 (YYYY-MM-DD)",
        )
        end = gr.Textbox(
            value=dt.date.today().isoformat(),
            label="종료일 (YYYY-MM-DD)",
        )

    with gr.Row():
        model_id = gr.Dropdown(
            choices=[
                "amazon/chronos-bolt-tiny",
                "amazon/chronos-bolt-mini",
                "amazon/chronos-bolt-small",
                "amazon/chronos-bolt-base",
            ],
            value="amazon/chronos-bolt-small",
            label="모델",
        )
        device = gr.Dropdown(
            choices=["cpu"],  # 필요하면 "cuda" 추가
            value="cpu",
            label="디바이스",
        )
        interval = gr.Dropdown(
            choices=["1d", "1h", "30m", "15m", "5m"],
            value="1d",
            label="간격 (Binance interval)",
        )

    btn = gr.Button("예측 실행")

    plot = gr.Plot(label="History + Forecast")
    table = gr.Dataframe(label="예측 결과 (분위수)")
    note = gr.Markdown()

    btn.click(
        fn=run_forecast,
        inputs=[ticker, start, end, horizon, model_id, device, interval],
        outputs=[plot, table, note],
    )

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