"""Grid search for a geometric adstock rate on the Vexpower practice dataset.

Download mmm-practice-synthetic.csv from vexpower.com/downloads/mmm/ into the
same folder, then run: python adstock-grid-search.py

Read this first: the practice CSV is synthetic and exactly linear, with no
noise, no adstock and no saturation (revenue = 20000 + 1.4*search +
0.8*social + 0.6*video + 3000*sin(pi*t/26) + 1800*promotion). So a grid
search on the raw file can only ever find zero carryover, and that tells you
nothing about real video advertising. The script shows that first, then
plants a known carryover (video retention 0.5) plus noise and checks how well
a grid search recovers it, on one draw and then on ten different noise draws.
None of the fitted decay rates here are findings about any real channel.
From vexpower.com/mmm/adstock.
"""
import numpy as np
import pandas as pd

df = pd.read_csv("mmm-practice-synthetic.csv")  # vexpower.com/downloads/mmm/
t = np.arange(len(df))
video = df["video_spend"].to_numpy(dtype=float)


def geometric_adstock(x, theta):
    """Unnormalised: a_t = x_t + theta * a_(t-1)."""
    out = np.zeros(len(x))
    for i, v in enumerate(x):
        out[i] = v + (theta * out[i - 1] if i else 0.0)
    return out


THETAS = np.round(np.arange(0.0, 0.91, 0.1), 1)


def grid(y):
    rows = []
    for theta in THETAS:
        X = np.column_stack([
            np.ones(len(df)),
            np.sin(2 * np.pi * t / 52), np.cos(2 * np.pi * t / 52),  # yearly season
            df["promotion"], df["search_spend"], df["social_spend"],
            geometric_adstock(video, theta),
        ])
        coef, *_ = np.linalg.lstsq(X, y, rcond=None)
        sse = np.sum((y - X @ coef) ** 2)
        rows.append({"theta": theta, "sse": sse, "video_coef": coef[-1],
                     "video_contribution": coef[-1] * X[:, -1].sum()})
    return pd.DataFrame(rows)


# 1. The raw practice data: no carryover was ever put in, so theta = 0 wins
#    with an almost perfect fit. This is a property of the synthetic file.
print("1. Raw practice CSV (no carryover, no noise in the data):")
print(grid(df["revenue"].to_numpy(dtype=float)).round(3).to_string(index=False))

# 2. Plant some carryover we can try to recover: true video retention 0.5
#    (true coefficient 0.6 on the adstocked series), plus noise.
planted = geometric_adstock(video, 0.5)
truth = 0.6 * planted.sum()
rng = np.random.default_rng(7)
y = (df["revenue"] + 0.6 * (planted - video)
     + rng.normal(0, 600, len(df))).to_numpy()
print(f"\n2. Planted retention 0.5, noise sd 600, seed 7 "
      f"(planted video contribution {truth:,.0f}):")
print(grid(y).round(3).to_string(index=False))

# 3. The same experiment on ten noise draws: how much does the 'best' theta move?
print("\n3. Best theta and contribution at the true theta across ten noise draws:")
for seed in range(10):
    rng = np.random.default_rng(seed)
    y = (df["revenue"] + 0.6 * (planted - video)
         + rng.normal(0, 600, len(df))).to_numpy()
    res = grid(y)
    best = res.loc[res["sse"].idxmin()]
    at_true = res.loc[res["theta"] == 0.5, "video_contribution"].iloc[0]
    print(f"  seed {seed}: best theta {best['theta']:.1f}, "
          f"contribution at theta 0.5 = {at_true:,.0f}")