#!/usr/bin/env python3
"""
Script 04: Difference-in-Differences (Python)
Atlas: Causality Atlas — Ch 09 (Econometric / Quasi-Experimental)
Demonstrates basic and staggered DiD.
"""
import numpy as np
import pandas as pd
import statsmodels.api as sm
import statsmodels.formula.api as smf

print("=== Difference-in-Differences ===\n")
np.random.seed(42)
N, T = 200, 5

# ---- Basic 2×2 DiD ----
print("--- Basic 2×2 DiD ---")
N_half = N // 2

# Pre-treatment
Y_pre_treat = 5 + np.random.normal(0, 1, N_half)
Y_pre_control = 5 + np.random.normal(0, 1, N_half)
# Post-treatment (true effect = 2)
Y_post_treat = 7 + np.random.normal(0, 1, N_half)
Y_post_control = 5 + np.random.normal(0, 1, N_half)

df_2x2 = pd.DataFrame({
    'Y': np.concatenate([Y_pre_treat, Y_pre_control, Y_post_treat, Y_post_control]),
    'treat': np.concatenate([np.ones(N_half), np.zeros(N_half)] * 2),
    'post': np.concatenate([np.zeros(N_half*2), np.ones(N_half*2)]),
    'id': np.tile(np.arange(N_half*2), 2)
})
df_2x2['did'] = df_2x2['treat'] * df_2x2['post']

did_model = smf.ols('Y ~ treat + post + did', data=df_2x2).fit()
print(f"DiD estimate (interaction) = {did_model.params['did']:.3f}")
print(f"True effect = 2.0")

# ---- Staggered DiD ----
print("\n--- Staggered DiD ---")
# Generate staggered treatment timing
units = []
unit_id = 0
for _ in range(N):
    first_treat = np.random.choice([2, 3, 4, None], p=[0.3, 0.3, 0.2, 0.2])
    unit_effect = np.random.normal(0, 0.5)
    for t in range(1, T+1):
        treat = 1 if (first_treat is not None and t >= first_treat) else 0
        Y = unit_effect + 0.5 * t + 2 * treat + np.random.normal(0, 1)
        units.append({'id': unit_id, 'year': 2019 + t, 'Y': Y,
                      'first_treat': first_treat if first_treat else 0,
                      'treat': treat})
    unit_id += 1

df = pd.DataFrame(units)

# Two-way fixed effects DiD (can be biased with staggered adoption)
twfe = smf.ols('Y ~ treat + C(id) + C(year)', data=df).fit()
print(f"TWFE (potentially biased with staggered adoption) = {twfe.params['treat']:.3f}")

# For proper staggered DiD with recent methods, use R `did` package
print("\nNote: Staggered DiD with Callaway & Sant'Anna estimator requires")
print("the R `did` package (script 04_did.R). Python implementation is")
print("available through the `staggered` or `qte` packages.")

print("\n=== Script complete ===")