"""Offline forecast from a downloaded local MLX bundle."""
import argparse
import numpy as np
from timestral_mlx import TimestralForecaster

p = argparse.ArgumentParser(description=__doc__)
p.add_argument("checkpoint", help="Local checkpoint directory, e.g. ./5M")
p.add_argument("--horizon", type=int, default=96)
a = p.parse_args()
forecaster = TimestralForecaster.from_pretrained(a.checkpoint)
t = np.arange(512, dtype=np.float32)
context = 10 + np.sin(2 * np.pi * t / 24) + 0.005 * t
context[100:108] = np.nan
forecast = forecaster.forecast(context, prediction_length=a.horizon)
print("Quantiles:", forecast.quantiles.shape)
print("Levels:", forecast.quantile_levels)
print("First 10 median predictions:", forecast.median[0, 0, :10])