import numpy as np
import pandas as pd
from ergodic import forecast, timeseries
from ergodic.forecast import schedule, window_graph

# 1. Simulate 60 months. Sales follow last month's web visits.
rng = np.random.default_rng(0)
T = 60
month = pd.date_range("2021-01-31", periods=T, freq="ME")
visits = 100 + 1.5 * np.arange(T) + rng.normal(0, 4, T)
sales = 20 + 0.8 * np.r_[visits[0], visits[:-1]] + rng.normal(0, 2, T)
frame = pd.DataFrame({"month": month, "visits": visits, "sales": sales})
data = timeseries(frame, time="month", freq="ME")

# 2. Say which series drives which: visits one month earlier drive sales.
graph = window_graph([("visits", "sales", 1)])
print(graph)
print(schedule(data, graph, fit="staged"))

# 3. Forecast six months. Visits are forecast too, and their uncertainty reaches sales.
result = forecast(data, 6, graph=graph, fit="staged", draws=500, warmup=500, seed=0)
print(result)
print(result.to_frame().round(1).to_string(index=False))

# 4. Check the sampler before you trust the bands.
for record in result.diagnostics.values():
    print(record)

# 5. Take the sales forecast apart: baseline, season and driver.
waterfall = result.explain(nodes="sales")["sales"]
bars = waterfall.to_frame()[["factor", "value", "lo", "hi"]]
print(bars.round(1).to_string(index=False))
print(f"total {waterfall.total:.1f}; bars add up: {waterfall.relative_closure_error < 1e-9}")