from ergodic import interrupted_time_series, synthetic_control, timeseries
from ergodic.datasets import load_smoking

# 1. Load the panel: packs per person, by state and year.
smoking = load_smoking()
print(smoking.data.n_entities, "states,", smoking.data.periods[0], "to", smoking.data.periods[-1])

# 2. Fit a synthetic California from the other states, before 1989.
fit = synthetic_control(smoking.data, "cigsale", treated="California", at=1989)
for state, weight in sorted(fit.weights.items(), key=lambda item: -item[1])[:5]:
    print(f"  {state:12} {weight:.2f}")
print(f"pre-1989 tracking error: {fit.pre_rmspe:.1f} packs")

# 3. Read the effect: the mean gap after 1989, with placebo inference.
low, high = fit.interval()
print(f"effect {fit.value:.1f} packs, placebo band [{low:.1f}, {high:.1f}]")
print(f"placebo p-value {fit.p_value:.3f}, gap in {fit.periods[-1]}: {fit.gaps[-1]:.1f}")

# 4. Fit an interrupted time series on California alone.
rows = smoking.data.to_pandas()
california = rows[rows["state"] == "California"][["year", "cigsale"]]
its = interrupted_time_series(timeseries(california, time="year"), "cigsale", at=1989)
low, high = its.interval()
print(f"interrupted time series: {its.value:.1f} packs, 95% interval [{low:.1f}, {high:.1f}]")