"""
Confounding simulation: when seasonal demand drives both spend and sales.

From the Vexpower MMM library: https://vexpower.com/mmm/causal-inference-for-marketers

A hidden demand signal raises sales on its own, and the marketing team also
spends more when demand is high. A regression of sales on spend alone then
gives advertising credit for demand. The script compares four estimates of the
return on $1 of spend with the known truth:

  1. naive regression (sales ~ spend)
  2. regression controlling for the true demand signal
  3. regression controlling for a noisy season index (what you usually have)
  4. a randomised schedule (spend set independently of demand)

Requires numpy only. Run:  python confounding-simulation.py
Change the settings below and run it again. All data is synthetic.
"""
import numpy as np

# ---- Settings you can change -------------------------------------------
WEEKS = 104
TRUE_EFFECT = 1.5        # true sales per $1 of spend
FOLLOW = 0.7             # 0 = spend ignores demand, 0.95 = spend tracks demand closely
SEASONAL_SWING = 8000    # how much one standard deviation of demand moves weekly sales
INDEX_NOISE = 0.5        # measurement error in the season index you control for (0 = perfect)
SIMULATIONS = 500        # repeat the world this many times to see the spread
SEED = 2026
# ---------------------------------------------------------------------------


def ols_slope(y, *columns):
    """Coefficient on the first column after an intercept and any controls."""
    X = np.column_stack([np.ones(len(y)), *columns])
    coef, *_ = np.linalg.lstsq(X, y, rcond=None)
    return coef[1]


def one_world(rng, follow):
    t = np.arange(WEEKS)
    demand = np.sin(2 * np.pi * (t - 8) / 52) + 0.35 * rng.standard_normal(WEEKS)
    demand = (demand - demand.mean()) / demand.std()
    spend = 10000 + 3000 * (follow * demand + np.sqrt(1 - follow**2) * rng.standard_normal(WEEKS))
    spend = np.maximum(spend, 500)
    sales = 40000 + SEASONAL_SWING * demand + TRUE_EFFECT * spend + rng.normal(0, 2500, WEEKS)
    season_index = demand + INDEX_NOISE * rng.standard_normal(WEEKS)
    return spend, sales, demand, season_index


def main():
    rng = np.random.default_rng(SEED)
    results = {"naive": [], "true demand control": [], "noisy season control": [], "randomised spend": []}
    for _ in range(SIMULATIONS):
        spend, sales, demand, index = one_world(rng, FOLLOW)
        results["naive"].append(ols_slope(sales, spend))
        results["true demand control"].append(ols_slope(sales, spend, demand))
        results["noisy season control"].append(ols_slope(sales, spend, index))
        r_spend, r_sales, _, _ = one_world(rng, 0.0)
        results["randomised spend"].append(ols_slope(r_sales, r_spend))

    print(f"True effect: ${TRUE_EFFECT:.2f} of sales per $1 of spend")
    print(f"Spend follows demand: {FOLLOW:.0%}, season index noise: {INDEX_NOISE}\n")
    print(f"{'Estimate':<24}{'mean':>8}{'5th pct':>10}{'95th pct':>10}{'bias':>9}")
    for name, vals in results.items():
        v = np.array(vals)
        print(f"{name:<24}{v.mean():>8.2f}{np.percentile(v, 5):>10.2f}{np.percentile(v, 95):>10.2f}{v.mean() - TRUE_EFFECT:>+9.2f}")
    print("\nThe naive estimate is biased by roughly SEASONAL_SWING x FOLLOW / 3000 (the spend scale).")


if __name__ == "__main__":
    main()