import numpy as np
import pandas as pd
from ergodic import forecast, hts
from ergodic.forecast import reconcile

# 1. Simulate 48 months of sales in two regions; the total is their sum.
rng = np.random.default_rng(0)
T = 48
month = pd.date_range("2022-01-31", periods=T, freq="ME")
north = 60 + 1.0 * np.arange(T) + rng.normal(0, 2, T)
south = 40 + 0.5 * np.arange(T) + rng.normal(0, 2, T)
frame = pd.DataFrame({"month": month, "total": north + south, "north": north, "south": south})
data = hts(frame, time="month", freq="ME", hierarchy={"total": ["north", "south"]})
tree = data.hierarchy
print(tree)

# 2. Forecast every level with its own model, and leave the levels as they are.
raw = forecast(data, 6, reconcile="none", draws=400, warmup=400, seed=0)
print(raw)

# 3. The levels disagree: the total's forecast is not the sum of the regions'.
gap = raw.mean("total") - raw.mean("north") - raw.mean("south")
print("total minus regions:", gap.round(1))

# 4. Reconcile four ways. Each reuses the same draws; nothing is refitted.
recipes = [
    ("bottom_up", {}),
    ("top_down", {"data": data}),
    ("mint", {}),
    ("mint", {"weights": "mint_shrink"}),
]
for method, options in recipes:
    result = reconcile(raw, tree, method=method, **options)
    final = result.mean("total")[-1]
    print(f"{result.reconciliation:32} coherent={result.coherent}  total in month 6: {final:.1f}")