import numpy as np
import pandas as pd
from ergodic import did, event_study, panel

# 1. Simulate 120 stores over 2010-2019. Some start a promotion in 2014, some in
#    2017, the rest never. The effect is 1.0 in the first year and grows 0.5 a year.
rng = np.random.default_rng(0)
rows = []
for store in range(120):
    start = int(rng.choice([2014, 2017, 0]))  # 0: never starts
    level = rng.normal(0, 1)
    for year in range(2010, 2020):
        promo = start > 0 and year >= start
        effect = 1.0 + 0.5 * (year - start) if promo else 0.0
        sales = level + 0.3 * (year - 2010) + effect + rng.normal(0, 0.5)
        rows.append({"store": store, "year": year, "promo": promo, "sales": sales})

# 2. Build the panel: one row per store and year.
data = panel(pd.DataFrame(rows), entity="store", time="year")
print(data.n_entities, "stores,", data.n_periods, "years, balanced:", data.is_balanced)

# 3. Estimate the average effect on the stores that ran the promotion.
result = did(data, "sales", "promo", random_state=0)
low, high = result.interval()
print(f"{result.flavor}: ATT {result.value:.2f}, 95% interval [{low:.2f}, {high:.2f}]")

# 4. Compare with the two-way fixed-effects regression.
twfe = did(data, "sales", "promo", strategy="twfe", random_state=0)
print(f"twfe: ATT {twfe.value:.2f}")

# 5. Trace the effect by years since the promotion started.
study = event_study(data, "sales", "promo", leads=3, random_state=0)
for row in study.by_event_time:
    low, high = row.interval
    print(f"{row.event_time:+d}  {row.estimate:5.2f}  [{low:5.2f}, {high:5.2f}]")