"""
Kronos Filter — HuggingFace Space (ZeroGPU)

Gradio app that wraps Kronos inference:
1. Visual UI (this page)
2. REST API at /call/predict (consumed by kronos-filter.vercel.app)

Runs on HF Spaces free CPU tier. @spaces.GPU decorator is kept as a no-op
so the code works on both CPU Spaces and (if upgraded to ZeroGPU via Pro)
GPU Spaces.
"""
from __future__ import annotations

import logging
import sys
from dataclasses import asdict
from datetime import datetime, timezone
from pathlib import Path

# Kronos model code lives at ./model/ in this Space
SPACE_ROOT = Path(__file__).resolve().parent
if str(SPACE_ROOT) not in sys.path:
    sys.path.insert(0, str(SPACE_ROOT))

import gradio as gr
import numpy as np
import torch

# @spaces.GPU decorator only useful on ZeroGPU (requires HF Pro).
# On free CPU tier we import it conditionally as a no-op.
try:
    import spaces
    SPACES_AVAILABLE = True
except ImportError:
    SPACES_AVAILABLE = False
    class _NoopGPU:
        def __call__(self, *a, **kw):
            def decorator(fn):
                return fn
            if len(a) == 1 and callable(a[0]):
                return a[0]
            return decorator
    class _NoopSpaces:
        GPU = _NoopGPU()
    spaces = _NoopSpaces()

from fetcher import fetch_aligned, build_future_timestamps
from predictor import KronosPathsPredictor
from probability_analyzer import analyze_ticker

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

# Predictor is lazily initialized on first GPU call, then cached
PREDICTOR: KronosPathsPredictor | None = None


def _ensure_predictor_on_device(device: str) -> KronosPathsPredictor:
    """
    Lazy-init predictor and ensure it's on the target device.
    With ZeroGPU, GPU is only accessible inside @spaces.GPU blocks, so we
    initialize on CPU first then migrate to GPU inside the decorator.
    """
    global PREDICTOR
    if PREDICTOR is None:
        logger.info(f"Initializing Kronos-small on {device}...")
        PREDICTOR = KronosPathsPredictor(model_size="small", device=device)
    elif PREDICTOR.device != device:
        logger.info(f"Moving Kronos from {PREDICTOR.device} to {device}...")
        PREDICTOR.model = PREDICTOR.model.to(device)
        PREDICTOR.tokenizer = PREDICTOR.tokenizer.to(device)
        PREDICTOR.device = device
    return PREDICTOR


@spaces.GPU(duration=180)
def run_prediction(
    tickers_raw: str,
    lookback: int,
    pred_len: int,
    samples: int,
    seed_raw: str,
) -> dict:
    """
    Main prediction pipeline — runs on ZeroGPU (A10G when available).
    Duration budget: 180s max per call.
    """
    # Parse tickers (comma or newline separated)
    tickers = []
    for chunk in tickers_raw.replace("\n", ",").split(","):
        t = chunk.strip().upper()
        if t and t not in tickers:
            tickers.append(t)
    if not tickers:
        raise gr.Error("Please provide at least one ticker.")
    if len(tickers) > 30:
        raise gr.Error("Max 30 tickers per request.")

    lookback = int(lookback)
    pred_len = int(pred_len)
    samples = int(samples)
    seed = None
    if seed_raw and seed_raw.strip().isdigit():
        seed = int(seed_raw.strip())

    logger.info(f"Request: tickers={tickers}, lookback={lookback}, pred_len={pred_len}, samples={samples}, seed={seed}")

    # 1. Fetch OHLCV (on CPU, no GPU needed)
    ok_list, failed = fetch_aligned(tickers, lookback=lookback, interval="1d")
    if not ok_list:
        raise gr.Error(f"Could not fetch any valid ticker data. Failed: {failed}")

    # 2. Future timestamps
    future_ts_list = [
        build_future_timestamps(td.timestamps.iloc[-1], pred_len, "1d")
        for td in ok_list
    ]

    # 3. Ensure model is on GPU (ZeroGPU allocates cuda:0 inside this decorator)
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    logger.info(f"Inference device: {device} (cuda available: {torch.cuda.is_available()})")
    predictor = _ensure_predictor_on_device(device)

    paths_list = predictor.predict_paths(
        df_list=[td.df for td in ok_list],
        x_timestamp_list=[td.timestamps for td in ok_list],
        y_timestamp_list=future_ts_list,
        pred_len=pred_len,
        T=1.0, top_p=0.9,
        sample_count=samples,
        verbose=False,
        seed=seed,
    )

    # 4. Compute metrics (CPU, fast)
    metrics_list = []
    for td, paths in zip(ok_list, paths_list):
        try:
            m = analyze_ticker(td.ticker, paths, initial_close=float(td.df["close"].iloc[-1]))
            metrics_list.append(m)
        except Exception as exc:
            logger.warning(f"[{td.ticker}] analysis failed: {exc}")

    # 5. Build JSON report (same shape as local CLI's latest.json)
    report = {
        "generated_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "config": {
            "lookback": lookback,
            "pred_len": pred_len,
            "samples": samples,
            "model_size": "small",
            "interval": "1d",
            "T": 1.0,
            "top_p": 0.9,
            "seed": seed,
            "device": device,
        },
        "failed_tickers": failed,
        "data_last_date": ok_list[0].timestamps.iloc[-1].isoformat() if ok_list else None,
        "tickers": sorted(
            [asdict(m) for m in metrics_list],
            key=lambda x: x["prob_up"],
            reverse=True,
        ),
    }

    # Path summaries for fan chart
    path_summaries = []
    for td, paths, ts_future in zip(ok_list, paths_list, future_ts_list):
        close_paths = paths[:, :, 3]
        path_summaries.append({
            "ticker": td.ticker,
            "history_dates": [d.isoformat() for d in td.timestamps.iloc[-60:]],
            "history_close": td.df["close"].iloc[-60:].tolist(),
            "future_dates": [d.isoformat() for d in ts_future],
            "median_path": np.median(close_paths, axis=0).tolist(),
            "p5_path": np.percentile(close_paths, 5, axis=0).tolist(),
            "p25_path": np.percentile(close_paths, 25, axis=0).tolist(),
            "p75_path": np.percentile(close_paths, 75, axis=0).tolist(),
            "p95_path": np.percentile(close_paths, 95, axis=0).tolist(),
        })
    report["path_summaries"] = path_summaries

    return report


# ----------------------------------------------------------------- UI

with gr.Blocks(title="Kronos Filter API") as demo:
    gr.Markdown(
        """
        # 📈 Kronos Filter — Probability API (ZeroGPU)

        Compute backend for [kronos-filter.vercel.app](https://kronos-filter.vercel.app).
        Fill in tickers to get probabilistic forecasts, or call programmatically via `/call/predict`.

        > ⚙️ **Hardware**: ZeroGPU (A10G, dynamically allocated).
        > ⚠️ Not financial advice. Results are **model-subjective probabilities**.
        """
    )
    with gr.Row():
        with gr.Column(scale=2):
            tickers_in = gr.Textbox(
                label="Tickers (comma or newline separated)",
                placeholder="MSFT, NVDA, AAPL",
                lines=3,
                value="MSFT, NVDA, AAPL",
            )
            with gr.Row():
                lookback_in = gr.Number(label="Lookback", value=400, minimum=50, maximum=512, precision=0)
                pred_len_in = gr.Number(label="Pred len", value=30, minimum=1, maximum=120, precision=0)
                samples_in = gr.Number(label="Samples", value=30, minimum=5, maximum=60, precision=0)
                seed_in = gr.Textbox(label="Seed (blank = random)", value="42")
            run_btn = gr.Button("🚀 Predict", variant="primary")
        with gr.Column(scale=3):
            output_json = gr.JSON(label="Probability report", open=False)

    run_btn.click(
        fn=run_prediction,
        inputs=[tickers_in, lookback_in, pred_len_in, samples_in, seed_in],
        outputs=output_json,
        api_name="predict",
    )

    gr.Markdown(
        """
        ---
        **Model**: Kronos-small (24.7M params, 512 context) · ZeroGPU A10G ·
        First request cold-start ~15s (model load onto GPU), subsequent calls ~5-15s.
        Daily GPU quota for free HF accounts: ~5-10 min.
        Source: [github.com/az30667176-sudo/kronos-filter](https://github.com/az30667176-sudo/kronos-filter)
        """
    )


demo.queue(default_concurrency_limit=1)

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