import numpy as np
import pandas as pd
from ergodic import backtest, timeseries
from ergodic.forecast import Naive, SeasonalNaive
from ergodic.welt import WeltForecaster

# 1. Simulate 120 days of history and the 14 after, with a promotion every Thursday.
rng = np.random.default_rng(0)
days = pd.date_range("2026-01-01", periods=134, freq="D")
promo = (days.dayofweek == 3).astype(float)
sales = 200 + 0.5 * np.arange(134) + 30 * promo + rng.normal(0, 5, 134)
frame = pd.DataFrame({"day": days, "sales": sales, "promo": promo})
history, ahead = frame.iloc[:120], frame.iloc[120:]
data = timeseries(history, time="day", freq="D", types={"promo": "continuous"})

# 2. Forecast 14 days, telling Chronos-2 the promotion calendar: one paid job.
model = WeltForecaster(future_covariates=["promo"])
result = model.fit(data).forecast(14, future=ahead[["day", "promo"]], seed=0)
print(result)
print(result.median().round(0))
record = result.diagnostics["sales"]
print(record.model_version, record.configuration, len(record.quantile_levels))

# 3. Backtest it at three origins against two free baselines: three paid jobs.
models = [WeltForecaster(future_covariates=["promo"]), Naive(), SeasonalNaive(period=7)]
reports = [backtest(m, data, test_size=14, n_splits=3) for m in models]
columns = ["mae", "crps", "coverage"]
table = pd.DataFrame({r.model: r.summary().loc["sales", columns] for r in reports}).T
print(table.round(2).to_string())
print("paid forecasts:", models[0].forecasts_used)