import os
import datetime as dt
import pandas as pd
import torch
import gradio as gr
import yfinance as yf

from chronos import BaseChronosPipeline  # from 'chronos-forecasting'

# ---- 전역 캐시: 모델을 한 번만 로드해 재사용 ----
_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]

# ---- 주가/크립토 데이터 로딩 (yfinance, 견고화) ----
def load_close_series(ticker: str, start: str, end: str, interval: str = "1d"):
    """
    BTC-USD 등 크립토 심볼에서 간헐적으로 timezone/파싱 오류가 나므로
    history() 경로를 우선 사용하고, 실패 시 한 번 재시도.
    """
    # 기본값 보정: 너무 최근만 고르면 빈 데이터가 나올 수 있어 일봉은 과거부터 권장
    _start = start or "2014-09-17"  # BTC-USD 최초 상장일 근처
    _end = end or dt.date.today().isoformat()

    tk = yf.Ticker(ticker)
    try:
        df = tk.history(start=_start, end=_end, interval=interval, auto_adjust=True, actions=False)
        if df.empty or "Close" not in df:
            raise ValueError("empty")
    except Exception:
        # fallback: download() 경로 시도
        df = yf.download(ticker, start=_start, end=_end, interval=interval, progress=False, threads=False)
        if df.empty or "Close" not in df:
            raise ValueError("데이터가 없거나 'Close' 열이 없습니다. 티커/날짜/간격을 확인하세요.")

    s = df["Close"].dropna().astype(float)
    if s.empty:
        raise ValueError("다운로드 결과가 비어 있습니다. 기간/간격을 줄이거나 다시 시도하세요.")
    return s

# ---- 예측 함수 (Gradio가 호출) ----
def run_forecast(ticker, start_date, end_date, horizon, model_id, device, interval):
    try:
        series = load_close_series(ticker.strip(), start_date, end_date, interval)
    except Exception as e:
        # Gradio v4에서는 Plot.update가 없음 → None 반환으로 정리
        return None, pd.DataFrame(), f"데이터 로딩 오류: {e}"

    pipe = get_pipeline(model_id, device)
    H = int(horizon)

    # Chronos 입력: 1D 텐서 (float)
    context = torch.tensor(series.values, dtype=torch.float32)

    # 예측: (num_series=1, num_quantiles=3, H) with q=[0.1, 0.5, 0.9]
    preds = pipe.predict(context=context, prediction_length=H)[0]
    q10, q50, q90 = preds[0], preds[1], preds[2]

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

    # 미래 x축: interval에 맞는 pandas 주기
    import matplotlib.pyplot as plt
    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:]

    # 그래프
    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 ({interval}, H={H})")
    plt.legend()
    plt.tight_layout()

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

# ---- Gradio UI ----
with gr.Blocks(title="Chronos Stock/Crypto Forecast") as demo:
    gr.Markdown("# Chronos 주가·크립토 예측 데모")
    with gr.Row():
        ticker = gr.Textbox(value="BTC-USD", label="티커 (예: AAPL, MSFT, 005930.KS, BTC-USD)")
        horizon = gr.Slider(5, 365, value=90, step=1, label="예측 스텝 H (간격 단위와 동일)")
    with gr.Row():
        start = gr.Textbox(value="2014-09-17", label="시작일 (YYYY-MM-DD, 예: 2014-09-17)")
        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"], value="cpu", label="디바이스")
        interval = gr.Dropdown(
            choices=["1d", "1h", "30m", "15m", "5m"],
            value="1d",
            label="간격"
        )
    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()