#!/usr/bin/env python3
"""Eldric Nexus F1 — complete runnable example.

    pip install https://repo.eldric.ai/models/eldric_nexus_f1-1.0.0-py3-none-any.whl
    python example_forecast.py

Downloads the weights on first run (315 MB, cached afterwards), forecasts 96
steps of a synthetic hourly series, and scores the result against two naive
baselines so you can see whether it is doing anything.

No GPU needed. Takes about a second on CPU once the weights are cached.
"""

import numpy as np
import torch

from eldric_nexus_f1 import TimeseriesType, load_model

CONTEXT, HORIZON = 512, 96
QUANTILE_LEVELS = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]


def make_series(n: int, seed: int = 7) -> np.ndarray:
    """Hourly data: level + trend + daily season + noise. float32 on purpose."""
    t = np.arange(n, dtype=np.float32)
    rng = np.random.default_rng(seed)
    x = 100.0 + 0.05 * t + 12.0 * np.sin(2 * np.pi * t / 24.0)
    return (x + 2.0 * rng.standard_normal(n)).astype(np.float32)


def main() -> None:
    full = make_series(CONTEXT + HORIZON)
    history, truth = full[:CONTEXT], full[CONTEXT:]

    # ---- 1. load the model ---------------------------------------------------
    # device="cpu" is the default. Use "cuda" or "mps" if you have them.
    model = load_model()

    # ---- 2. wrap the history ------------------------------------------------
    # target is 2D: [variates, context]. A univariate series is still 2D.
    # It MUST be float32 -- float64 (numpy/pandas default) fails inside PyTorch.
    series = TimeseriesType(
        target=torch.from_numpy(history).unsqueeze(0),
        past_covariates=None,
        future_covariates=None,
    )

    # ---- 3. forecast --------------------------------------------------------
    # Returns a LIST, one entry per input series, each (variates, quantiles, horizon).
    forecasts = model.forecast([series], prediction_length=HORIZON, output_type="numpy")
    q = np.asarray(forecasts[0])[0]                     # (9, HORIZON)
    median = q[QUANTILE_LEVELS.index(0.5)]              # index 4

    # ---- 4. look at it ------------------------------------------------------
    print(f"history {CONTEXT} steps -> forecast {HORIZON} steps, shape {q.shape}")
    print()
    print("  step        q10        q50        q90       truth")
    for i in list(range(3)) + [HORIZON // 2, HORIZON - 1]:
        print(f"  {i:4d}  {q[0, i]:9.3f}  {median[i]:9.3f}  {q[8, i]:9.3f}  {truth[i]:10.3f}")
    print()

    mae_model = np.abs(median - truth).mean()
    mae_carry = np.abs(history[-1] - truth).mean()        # repeat the last value
    mae_season = np.abs(history[-24:][: min(24, HORIZON)].mean() - truth).mean()
    print(f"  MAE of the median            {mae_model:8.3f}")
    print(f"  MAE of repeat-last-value     {mae_carry:8.3f}   "
          f"({mae_model / mae_carry:.2f}x — under 1.00 is better)")
    print(f"  MAE of last-day mean         {mae_season:8.3f}")
    print()
    inside = np.mean((truth >= q[0]) & (truth <= q[8]))
    print(f"  truth inside the q10..q90 band: {100 * inside:.1f} %")


if __name__ == "__main__":
    main()