"""
Geo test analysis: three answers from one test.

Reproduces the numbers on https://vexpower.com/mmm/geo-lift-testing using
geo-experiment-synthetic.csv (Vexpower synthetic teaching data: 12 regions,
6 matched pairs, 4 pre weeks and 4 post weeks).

    python geo-test-analysis.py geo-experiment-synthetic.csv

Needs pandas and numpy. Works on your own file if it has the columns
region, arm (treatment/comparison), period (pre/post), week and revenue_k;
pair, media_spend_k and known_media_lift_k are optional.
"""
import sys

import numpy as np
import pandas as pd

path = sys.argv[1] if len(sys.argv) > 1 else "geo-experiment-synthetic.csv"
df = pd.read_csv(path)
df["treated"] = (df["arm"].str.lower().str.startswith("t")).astype(int)
df["post"] = (df["period"].str.lower() == "post").astype(int)


def three_answers(d: pd.DataFrame, label: str) -> None:
    means = d.groupby(["arm", "period"])["revenue_k"].mean().unstack("period")[["pre", "post"]]
    means["change"] = means["post"] - means["pre"]
    before_after = means.loc["treatment", "change"]
    did = before_after - means.loc["comparison", "change"]
    print(f"\n=== {label} ===")
    print(means.round(2).to_string())
    print(f"Before/after, treatment only : {before_after:.2f}")
    print(f"Difference-in-differences    : {did:.2f}")
    t_post = d[(d.treated == 1) & (d.post == 1)]
    if "known_media_lift_k" in d:
        print(f"Known media effect           : {t_post['known_media_lift_k'].mean():.2f}")
    if "media_spend_k" in d:
        spend = t_post["media_spend_k"].mean()
        if spend > 0:
            print(f"iROAS (DiD / extra spend)    : {did / spend:.2f}")


three_answers(df, "All regions")

# The region with the local promotion is visible here only because the
# synthetic file includes the answer-key column. In a real test you would find
# it from field notes, a residual check or a pair-level comparison like below.
if "promotion_lift_k" in df:
    promo = df.loc[df["promotion_lift_k"] > 0, "region"].unique()
    three_answers(df[~df["region"].isin(promo)], f"Excluding {', '.join(promo)}")

# Pair-level DiD: an outlying pair is a prompt to go and ask what happened there.
if "pair" in df:
    print("\n=== DiD by matched pair ===")
    g = df.groupby(["pair", "arm", "period"])["revenue_k"].mean().unstack(["arm", "period"])
    pair_did = (g[("treatment", "post")] - g[("treatment", "pre")]) - (g[("comparison", "post")] - g[("comparison", "pre")])
    print(pair_did.round(2).to_string())

# The same DiD as a two-way fixed effects regression: region and week dummies
# plus treated x post. With one start date for every treated region this is
# the textbook 2x2 estimate.
regions = pd.get_dummies(df["region"], drop_first=True, dtype=float)
weeks = pd.get_dummies(df["week"], drop_first=True, dtype=float)
X = np.column_stack([np.ones(len(df)), (df.treated * df.post).values, regions.values, weeks.values])
beta, *_ = np.linalg.lstsq(X, df["revenue_k"].values, rcond=None)
print(f"\nTwo-way fixed effects estimate of treated x post: {beta[1]:.2f}")