#!/usr/bin/env python3
"""
Script 02: Propensity Score — Estimation, Matching, IPW (Python)
Atlas: Causality Atlas — Ch 05 (Neyman-Rubin Potential Outcomes)
"""
import numpy as np
import pandas as pd
from sklearn.linear_model import LogisticRegression
import statsmodels.api as sm
import statsmodels.formula.api as smf

print("=== Propensity Score Analysis ===\n")
np.random.seed(42)
N = 1000

# Simulate confounded data
C1, C2 = np.random.normal(0, 1, (2, N))
C3 = np.random.binomial(1, 0.5, N)

logit_ps = -1 + 0.5*C1 - 0.8*C2 + 0.6*C3
ps = 1 / (1 + np.exp(-logit_ps))
A = np.random.binomial(1, ps)
Y = 2 + 1.5*A + 0.3*C1 - 0.5*C2 + 0.4*C3 + np.random.normal(0, 1, N)

df = pd.DataFrame({'Y': Y, 'A': A, 'C1': C1, 'C2': C2, 'C3': C3})

# Naive estimate
naive = smf.ols('Y ~ A', data=df).fit()
print(f"Naive ATE = {naive.params['A']:.3f} (confounded)")

# ---- 1. PS Estimation ----
ps_model = LogisticRegression(C=1e6).fit(df[['C1','C2','C3']], df['A'])
df['ps'] = ps_model.predict_proba(df[['C1','C2','C3']])[:, 1]
print(f"PS range: [{df['ps'].min():.3f}, {df['ps'].max():.3f}]")

# ---- 2. IPW (ATE) ----
def compute_ipw(df):
    """Compute stabilized IPW weights."""
    ps = df['ps'].values
    p_treat = df['A'].mean()
    w = np.where(df['A'] == 1, p_treat/ps, (1-p_treat)/(1-ps))
    return w

df['w'] = compute_ipw(df)
ipw_model = smf.wls('Y ~ A', data=df, weights=df['w']).fit(cov_type='HC3')
print(f"ATE (IPW) = {ipw_model.params['A']:.3f} (SE={ipw_model.bse['A']:.3f})")

# ---- 3. Nearest-Neighbor Matching (simple 1:1) ----
print("\n--- Nearest-Neighbor Matching ---")
treated = df[df['A'] == 1].copy()
control = df[df['A'] == 0].copy()

matches = []
for i, t_row in treated.iterrows():
    # Find nearest control by PS
    control['dist'] = np.abs(control['ps'] - t_row['ps'])
    best = control.loc[control['dist'].idxmin()]
    matches.append({'treated_Y': t_row['Y'], 'control_Y': best['Y'],
                    'treated_ps': t_row['ps'], 'control_ps': best['ps']})
    control = control.drop(best.name)

matches_df = pd.DataFrame(matches)
att_est = (matches_df['treated_Y'] - matches_df['control_Y']).mean()
print(f"ATT (1:1 matching) = {att_est:.3f}")

# ---- 4. Doubly Robust (AIPW) ----
print("\n--- Doubly Robust (AIPW) ---")
outcome = smf.ols('Y ~ A + C1 + C2 + C3', data=df).fit()
df['mu1'] = outcome.predict(df.assign(A=1))
df['mu0'] = outcome.predict(df.assign(A=0))
# AIPW formula
aipw = ((df['A']*df['Y'] - (df['A'] - df['ps'])*df['mu1'])/df['ps'] -
        ((1-df['A'])*df['Y'] + (df['A'] - df['ps'])*df['mu0'])/(1 - df['ps'])).mean()
print(f"AIPW ATE = {aipw:.3f}")

# ---- 5. Balance check ----
print("\n--- Balance Check ---")
for c in ['C1', 'C2', 'C3']:
    smd_before = (df[df['A']==1][c].mean() - df[df['A']==0][c].mean()) / df[c].std()
    # Weighted SMD
    w_treated = df[df['A'] == 1]['w'].values
    w_control = df[df['A'] == 0]['w'].values
    w_mean_t = np.average(df[df['A']==1][c], weights=w_treated)
    w_mean_c = np.average(df[df['A']==0][c], weights=w_control)
    smd_after = (w_mean_t - w_mean_c) / np.sqrt(
        (np.cov(df[df['A']==1][c], aweights=w_treated) +
         np.cov(df[df['A']==0][c], aweights=w_control)) / 2
    )
    print(f"  {c}: SMD(before)={abs(smd_before):.3f}, SMD(after)={abs(smd_after):.3f}")

print(f"\nTrue ATE = 1.5")
print("=== Script complete ===")