import numpy as np
import pandas as pd
from ergodic import SCM, dag, tabular

# 1. Simulate 500 weeks. Demand raises both ad spend and sales.
rng = np.random.default_rng(1)
n = 500
demand = rng.normal(size=n)
ads = 0.5 * demand + rng.normal(size=n)
sales = 2.0 * ads + 3.0 * demand + rng.normal(size=n)
weeks = pd.DataFrame({"demand": demand, "ads": ads, "sales": sales})

# 2. Fit the SCM over the true graph.
graph = dag(["demand -> ads", "demand -> sales", "ads -> sales"])
scm = SCM(graph, tabular(weeks), draws=500, tune=500, chains=1, random_state=0).fit()

# 3. Observe: the distribution of sales as the data shows it.
print(scm.marginal("sales"))

# 4. Intervene: set ad spend to 1 for every week, then to 0.
high = scm.interventional("sales", do={"ads": 1.0})
low = scm.do({"ads": 0.0}).interventional("sales")
print(f"E[sales | do(ads=1)] - E[sales | do(ads=0)] = {high.value - low.value:.3f}")

# 5. Ask a counterfactual: what would week 0 have sold with ads at 1?
week0 = weeks.head(1)
print(week0.round(3).to_string(index=False))
what_if = scm.counterfactual(tabular(week0), do={"ads": 1.0}, target="sales")
print(what_if)

# 6. Individual effects: one more unit of ads, from each week's own level.
ite = scm.ite(tabular(weeks.head(3)), "ads", "sales", shift=1.0)
print(ite.framing, ite.values.round(3))