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

# 1. Simulate 48 months. Sales grow, and fall by 12 for each unit of price.
rng = np.random.default_rng(0)
T = 48
month = pd.date_range("2022-01-31", periods=T, freq="ME")
price = 10 + 0.05 * np.arange(T) + rng.normal(0, 0.3, T)
sales = 300 + 2.0 * np.arange(T) - 12.0 * price + rng.normal(0, 3, T)
frame = pd.DataFrame({"month": month, "price": price, "sales": sales})
data = timeseries(frame, time="month", freq="ME")

# 2. Price acts on sales in the same month. You set the price, so you pin its future.
graph = window_graph([("price", "sales", 0)])
print(schedule(data, graph, scenario={"price": "hold"}))

# 3. Forecast under two price plans: hold the last price, or raise it by 10%.
plans = {"hold": "hold", "raise 10%": round(1.1 * price[-1], 2)}
results = {}
for name, plan in plans.items():
    results[name] = forecast(
        data, 6, graph=graph, scenario={"price": plan}, draws=400, warmup=400, seed=0
    )
    print(f"{name:10} price {results[name].median('price')[0]:.2f}", end="  ")
    print("sales", results[name].median("sales").round(1))

# 4. Compare the plans draw by draw: the fit is the same, only the price differs.
change = results["raise 10%"].draws("sales") - results["hold"].draws("sales")
low, high = np.percentile(change, [5, 95], axis=0)
print(f"change in month 1: {change.mean(axis=0)[0]:.1f} (90% band {low[0]:.1f} to {high[0]:.1f})")