"""TinyCast: Probabilistic Zero-Shot Forecasting — Gradio demo.

TinyCast is a 146K-parameter attention-free time-series foundation model that
forecasts unseen series zero-shot and returns nine quantile forecasts.
https://huggingface.co/raws-labs/tinycast
"""

import os
import io
import time
from pathlib import Path

import numpy as np
import pandas as pd
import gradio as gr
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

import torch

# ---------------------------------------------------------------------------
# Model loading — at module scope, CPU (the model is 146K params, ~0.6 MB).
# ---------------------------------------------------------------------------
from huggingface_hub import hf_hub_download

HF_REPO = "raws-labs/tinycast"

# Download config.json and model.safetensors (config.json lands beside weights)
_hf_config = hf_hub_download(HF_REPO, "config.json")
WEIGHTS_PATH = hf_hub_download(HF_REPO, "model.safetensors")

# The tinycast package ships on GitHub; we vendor it into the Space repo.
import sys
sys.path.insert(0, str(Path(__file__).parent))

from tinycast import TinyCastPredictor


# ---------------------------------------------------------------------------
# Frequency → (freq_label, domain) mapping for the seasonal scale factor.
# The model needs a pandas frequency string and a domain for the scale factor.
# ---------------------------------------------------------------------------
FREQ_MAP = {
    "Hourly (H)": ("H", "Energy"),
    "Daily (D)": ("D", "Sales"),
    "Weekly (W)": ("W", "Sales"),
    "Monthly (M)": ("M", "Sales"),
    "Quarterly (Q)": ("Q", "Economic"),
    "4-hourly (4H)": ("4H", "Energy"),
    "10-minute (10T)": ("10T", "Energy"),
    "Minute (T)": ("T", "Energy"),
}


def parse_csv(file_obj):
    """Parse an uploaded CSV file into a numpy array of values."""
    if file_obj is None:
        return None, None
    try:
        df = pd.read_csv(file_obj.name)
    except Exception:
        df = pd.read_csv(file_obj)

    # Try to find a value column: look for common names, else take last column
    value_cols = [c for c in df.columns if c.lower() in ("value", "y", "target", "v", "measurement")]
    if value_cols:
        values = df[value_cols[0]].values
    else:
        values = df.iloc[:, -1].values

    values = pd.to_numeric(pd.Series(values), errors="coerce").values.astype("float32")
    return values, df


def parse_text_input(text):
    """Parse pasted numeric values (comma or newline separated)."""
    if not text or not text.strip():
        return None
    parts = text.replace(",", " ").replace("\n", " ").split()
    vals = []
    for p in parts:
        try:
            vals.append(float(p))
        except ValueError:
            pass
    if not vals:
        return None
    return np.array(vals, dtype="float32")


def forecast(
    csv_file,
    text_input,
    freq_choice,
    prediction_length,
    flip_invariance,
):
    """Run TinyCast zero-shot forecast and return a quantile plot."""
    # Get the time series values
    if csv_file is not None:
        values, _ = parse_csv(csv_file)
    else:
        values = parse_text_input(text_input)

    if values is None or len(values) < 10:
        return None, "Please provide at least 10 data points (via CSV upload or text input)."

    freq_str, domain = FREQ_MAP.get(freq_choice, ("H", "Energy"))

    # Build gluonts-style entry. pandas prefers lowercase freq for Period,
    # but the model's seasonal_scale_factor expects uppercase (e.g. "H", "D").
    _period_freq = freq_str.lower() if freq_str.upper() == freq_str else freq_str
    try:
        start = pd.Period("2020-01-01", freq=_period_freq)
    except Exception:
        start = pd.Period("2020-01-01", freq="D")

    entry = {"target": values, "start": start, "item_id": "input"}

    # Build the predictor
    t0 = time.perf_counter()
    predictor = TinyCastPredictor(
        prediction_length=int(prediction_length),
        checkpoint_path=WEIGHTS_PATH,
        freq=freq_str,
        domain=domain,
        device="cpu",
        force_flip_invariance=bool(flip_invariance),
    )
    forecasts = predictor.predict([entry])
    elapsed = time.perf_counter() - t0

    fc = forecasts[0]
    # forecast_array is (Q, prediction_length) with quantile levels ascending
    arr = fc.forecast_array  # (Q, pl)
    quantile_levels = [float(k) for k in fc.forecast_keys]
    pl = arr.shape[1]

    # Build the plot
    fig, ax = plt.subplots(figsize=(12, 5))

    # Plot context (last 200 points for readability)
    context_show = min(200, len(values))
    ctx_x = np.arange(-context_show, 0)
    ax.plot(ctx_x, values[-context_show:], color="black", linewidth=1.0, label="Context (observed)")

    # Plot forecast quantiles
    fut_x = np.arange(0, pl)
    q_colors = plt.cm.viridis(np.linspace(0.15, 0.85, len(quantile_levels)))

    # Fill between outer quantiles for the confidence band
    if len(quantile_levels) >= 2:
        ax.fill_between(
            fut_x, arr[0, :], arr[-1, :],
            color="steelblue", alpha=0.15, label=f"{int(float(quantile_levels[0])*100)}%–{int(float(quantile_levels[-1])*100)}% interval",
        )
    if len(quantile_levels) >= 4:
        mid_lo = len(quantile_levels) // 4
        mid_hi = 3 * len(quantile_levels) // 4
        ax.fill_between(
            fut_x, arr[mid_lo, :], arr[mid_hi, :],
            color="steelblue", alpha=0.25, label=f"{int(float(quantile_levels[mid_lo])*100)}%–{int(float(quantile_levels[mid_hi])*100)}% interval",
        )

    # Median line
    qm = len(quantile_levels) // 2
    ax.plot(fut_x, arr[qm, :], color="crimson", linewidth=2.0, label="Median forecast")

    # Outer quantile lines
    for i in [0, -1]:
        ax.plot(fut_x, arr[i, :], color=q_colors[i], linewidth=0.8, linestyle="--", alpha=0.5)

    ax.axvline(x=0, color="gray", linestyle=":", linewidth=0.8)
    ax.set_xlabel("Time (relative)")
    ax.set_ylabel("Value")
    ax.set_title(f"TinyCast {int(prediction_length)}-step probabilistic forecast ({freq_choice})")
    ax.legend(loc="upper left", fontsize=8)
    fig.tight_layout()

    # Summary text
    info = (
        f"Context: {len(values)} points | Frequency: {freq_str} | Domain: {domain} | "
        f"Horizon: {pl} steps | Quantiles: {len(quantile_levels)} | "
        f"Flip-invariance: {'on' if flip_invariance else 'off'} | "
        f"Inference: {elapsed:.2f}s"
    )

    plt.close(fig)
    return fig, info


# ---------------------------------------------------------------------------
# Build the Gradio UI
# ---------------------------------------------------------------------------
CSS = """
#main-col { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="TinyCast Forecaster") as demo:
    gr.Markdown(
        "# TinyCast: Probabilistic Zero-Shot Forecasting\n\n"
        "An attention-free, 146K-parameter time-series foundation model that forecasts "
        "unseen series zero-shot and returns nine quantile forecasts. "
        "[Model](https://huggingface.co/raws-labs/tinycast) · "
        "[Paper](https://arxiv.org/abs/2608.15767) · "
        "[Code](https://github.com/raws-labs/tinycast)"
    )

    with gr.Column(elem_id="main-col"):
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### Input")
                csv_input = gr.File(label="Upload CSV (value column or last column used)", file_types=[".csv"])
                gr.Markdown("**— or —**")
                text_input = gr.Textbox(
                    label="Paste values (comma or newline separated)",
                    placeholder="10.0, 12.3, 9.8, 11.1, ...",
                    lines=3,
                )
                freq_choice = gr.Dropdown(
                    choices=list(FREQ_MAP.keys()),
                    value="Hourly (H)",
                    label="Frequency / Domain",
                )
                prediction_length = gr.Slider(
                    minimum=1, maximum=512, value=48, step=1,
                    label="Forecast horizon (steps)",
                )
                flip_inv = gr.Checkbox(
                    value=True,
                    label="Flip-invariance symmetrization (improves accuracy)",
                )
                run_btn = gr.Button("Forecast", variant="primary")

            with gr.Column(scale=2):
                gr.Markdown("### Forecast")
                plot_output = gr.Plot(label="Probabilistic forecast")
                info_output = gr.Textbox(label="Summary", interactive=False)

        run_btn.click(
            fn=forecast,
            inputs=[csv_input, text_input, freq_choice, prediction_length, flip_inv],
            outputs=[plot_output, info_output],
            api_name="forecast",
        )

        gr.Markdown("### Examples")
        gr.Examples(
            examples=[
                ["example_hourly_energy.csv", None, "Hourly (H)", 48, True],
                ["example_daily_sales.csv", None, "Daily (D)", 30, True],
                ["example_monthly_temp.csv", None, "Monthly (M)", 24, True],
            ],
            inputs=[csv_input, text_input, freq_choice, prediction_length, flip_inv],
            outputs=[plot_output, info_output],
            fn=forecast,
            cache_examples=True,
            cache_mode="lazy",
        )

        gr.Markdown(
            "\n\n---\n"
            "TinyCast replaces self-attention with dilated causal convolutions and a "
            "zero-parameter normalized-periodogram phase prior. The model is 146,505 "
            "parameters (~0.6 MB) and runs on CPU. License: Apache-2.0."
        )


if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)