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

# 1. Simulate 90 days of sales in two regions; the total is their sum.
rng = np.random.default_rng(0)
days = pd.date_range("2026-01-01", periods=90, freq="D")
north = 120 + 0.4 * np.arange(90) + rng.normal(0, 4, 90)
south = 80 + 0.2 * np.arange(90) + rng.normal(0, 4, 90)
frame = pd.DataFrame({"day": days, "total": north + south, "north": north, "south": south})
data = hts(frame, time="day", freq="D", hierarchy={"total": ["north", "south"]})
tree = data.hierarchy

# 2. Forecast all three levels in one paid job, and leave them as they are.
raw = WeltForecaster().fit(data).forecast(14, seed=0)
print(raw)
gap = raw.mean("total") - raw.mean("north") - raw.mean("south")
print("total minus regions:", gap.round(1))

# 3. Reconcile two ways. Both reuse the same draws.
for method in ["bottom_up", "mint"]:
    fixed = reconcile(raw, tree, method=method)
    print(f"{fixed.reconciliation:24} coherent={fixed.coherent}")