import numpy as np
import pandas as pd
from ergodic import backtest, timeseries
from ergodic.forecast import Drift, Naive, SeasonalNaive, forecast_model

# 1. Simulate 60 months of sales with a trend and a yearly season.
rng = np.random.default_rng(0)
T = 60
month = pd.date_range("2021-01-31", periods=T, freq="ME")
season = 10 * np.sin(2 * np.pi * np.arange(T) / 12)
sales = 100 + 1.5 * np.arange(T) + season + rng.normal(0, 3, T)
data = timeseries(pd.DataFrame({"month": month, "sales": sales}), time="month", freq="ME")

# 2. Line up three baselines and the Bayesian forecaster.
models = [
    Naive(),
    SeasonalNaive(period=12),
    Drift(),
    forecast_model(name="Bayesian", draws=500, warmup=500),
]

# 3. Refit each model at two origins and score the next six months.
reports = [backtest(model, data, test_size=6, n_splits=2, m=12) for model in models]
print(reports[-1])

# 4. Compare the mean scores across folds.
columns = ["mae", "rmse", "mase", "crps", "coverage"]
table = pd.DataFrame({r.model: r.summary().loc["sales", columns] for r in reports}).T
print(table.round(2).to_string())

# 5. Check how the sampler behaved in each fold.
diagnostics = reports[-1].diagnostics()
print(diagnostics[["fold", "cutoff", "divergences", "min_ess"]].round(1).to_string(index=False))