"""
Difference-in-differences and an event-study check, from scratch.

Companion to https://vexpower.com/mmm/difference-in-differences. Runs on
geo-experiment-synthetic.csv (Vexpower synthetic teaching data) or any long
panel with columns: week, region, arm (treatment/comparison), period (pre/post)
and an outcome column.

    python did-event-study.py geo-experiment-synthetic.csv revenue_k

Needs pandas and numpy. Standard errors are clustered by region (CR1). With
only a dozen regions they are rough; treat them as a sense of scale.
"""
import sys

import numpy as np
import pandas as pd

path = sys.argv[1] if len(sys.argv) > 1 else "geo-experiment-synthetic.csv"
y_col = sys.argv[2] if len(sys.argv) > 2 else "revenue_k"
df = pd.read_csv(path).sort_values(["region", "week"]).reset_index(drop=True)
df["treated"] = df["arm"].str.lower().str.startswith("t").astype(int)
df["post"] = (df["period"].str.lower() == "post").astype(int)

# 1. The 2x2 calculation from four means.
m = df.groupby(["treated", "post"])[y_col].mean()
did_2x2 = (m[1, 1] - m[1, 0]) - (m[0, 1] - m[0, 0])
print("Group means:\n", m.unstack().rename(index={0: "control", 1: "treatment"}, columns={0: "pre", 1: "post"}).round(2))
print(f"\n2x2 difference-in-differences: {did_2x2:.3f}")


def ols_cluster(X: np.ndarray, y: np.ndarray, groups: np.ndarray):
    """OLS coefficients with cluster-robust (CR1) standard errors."""
    XtX_inv = np.linalg.pinv(X.T @ X)
    beta = XtX_inv @ X.T @ y
    resid = y - X @ beta
    meat = np.zeros((X.shape[1], X.shape[1]))
    for g in np.unique(groups):
        idx = groups == g
        s = X[idx].T @ resid[idx]
        meat += np.outer(s, s)
    n, k, G = len(y), X.shape[1], len(np.unique(groups))
    scale = G / (G - 1) * (n - 1) / (n - k)
    se = np.sqrt(np.clip(np.diag(scale * XtX_inv @ meat @ XtX_inv), 0, None))
    return beta, se


# 2. Two-way fixed effects: region and week dummies plus treated x post.
R = pd.get_dummies(df["region"], drop_first=True, dtype=float).values
W = pd.get_dummies(df["week"], drop_first=True, dtype=float).values
X = np.column_stack([np.ones(len(df)), df["treated"] * df["post"], R, W])
beta, se = ols_cluster(X, df[y_col].values.astype(float), df["region"].values)
print(f"Two-way fixed effects estimate: {beta[1]:.3f} (cluster SE {se[1]:.3f})")

# 3. Event study: one coefficient per week relative to launch, with the
# last pre-launch week as the reference. Pre-launch coefficients near zero
# are what parallel trends looks like; a slope in them is a warning.
weeks = sorted(df["week"].unique())
first_post = min(df.loc[df["post"] == 1, "week"])
rel = {w: weeks.index(w) - weeks.index(first_post) for w in weeks}
df["rel"] = df["week"].map(rel)
leads_lags = [r for r in sorted(df["rel"].unique()) if r != -1]
D = np.column_stack([(df["treated"] * (df["rel"] == r)).values for r in leads_lags])
X_es = np.column_stack([np.ones(len(df)), D, R, W])
b_es, se_es = ols_cluster(X_es, df[y_col].values.astype(float), df["region"].values)
print("\nEvent study (treated x week relative to launch, reference = -1):")
for i, r in enumerate(leads_lags):
    print(f"  {r:+d}: {b_es[1 + i]:7.3f}  (SE {se_es[1 + i]:.3f})")
pre = [b_es[1 + i] for i, r in enumerate(leads_lags) if r < 0]
print(f"\nLargest pre-launch coefficient in absolute value: {max(map(abs, pre)) if pre else float('nan'):.3f}")