# Marketing mix modeling in Python, from scratch (Vexpower)
# Guide: https://vexpower.com/mmm/marketing-mix-modeling-python
# Run: python mmm-python-tutorial.py  (needs numpy, pandas, scipy, matplotlib)
# The '# %%' markers make each step a runnable cell in VS Code or Jupyter.

# %% [markdown]
# # Marketing mix modeling in Python, from scratch
#
# A complete marketing mix model built with numpy, pandas and scipy (plus matplotlib for the charts),
# on Vexpower's synthetic practice dataset. Every number in the guide at
# https://vexpower.com/mmm/marketing-mix-modeling-python comes from running this file.
#
# The data is invented for teaching. Put `mmm-practice-synthetic.csv` in the same folder as this
# notebook, or leave `DATA_URL` as it is to download it.

# %% setup
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.optimize import minimize

DATA_URL = "https://vexpower.com/downloads/mmm/mmm-practice-synthetic.csv"
CHANNELS = ["search_spend", "social_spend", "video_spend"]

try:
    df = pd.read_csv("mmm-practice-synthetic.csv", parse_dates=["week"])
except FileNotFoundError:
    df = pd.read_csv(DATA_URL, parse_dates=["week"])

np.set_printoptions(suppress=True)
pd.set_option("display.width", 120)

# %% [markdown]
# ## 1. Load and audit the data
#
# Before any modeling: how many rows, any gaps in the weeks, any missing values, and what the
# spend and revenue look like.

# %% audit
print(df.shape)
print("Missing values:", int(df.isna().sum().sum()))
print("Week gaps:", df["week"].diff().dropna().value_counts().to_dict())
print(df["week"].min().date(), "to", df["week"].max().date())
print(df.drop(columns="week").describe().round(1).T[["mean", "min", "max"]])
print(df[["revenue"] + CHANNELS + ["promotion"]].corr().round(2)["revenue"])

# %% [markdown]
# ## 2. Adstock and saturation
#
# Geometric adstock carries a share of last week's advertising into this week. The Hill curve
# (with shape fixed at 1) bends the adstocked spend so that each extra dollar does a little less.

# %% transforms
def geometric_adstock(x, decay):
    """adstock[t] = x[t] + decay * adstock[t-1]"""
    out = np.zeros(len(x))
    carry = 0.0
    for t, value in enumerate(x):
        carry = value + decay * carry
        out[t] = carry
    return out


def hill(x, half_sat):
    """Share of the maximum effect reached at x. Equals 0.5 when x == half_sat."""
    return x / (x + half_sat)

# %% [markdown]
# ## 3. Build the design matrix and fit OLS for one set of parameters
#
# Controls: intercept, linear trend, one yearly sine/cosine pair for seasonality, and the
# promotion flag. Then one transformed column per channel. Half-saturation is expressed as a
# multiple of the channel's average weekly spend so the three channels share a scale.

# %% design
n = len(df)
t = np.arange(n)
controls = np.column_stack([
    np.ones(n),                   # intercept (baseline)
    t,                            # trend
    np.sin(2 * np.pi * t / 52),   # yearly seasonality
    np.cos(2 * np.pi * t / 52),
    df["promotion"].values,
])
control_names = ["intercept", "trend", "season_sin", "season_cos", "promotion"]
y = df["revenue"].values.astype(float)
spend = df[CHANNELS].values.astype(float)
mean_spend = spend.mean(axis=0)


def media_columns(params, spend=spend):
    decays, half_mults = params[:3], params[3:]
    cols = []
    for j in range(3):
        a = geometric_adstock(spend[:, j], decays[j])
        cols.append(hill(a, half_mults[j] * mean_spend[j]))
    return np.column_stack(cols)


def fit_ols(params, rows=slice(None)):
    X = np.column_stack([controls, media_columns(params)])
    beta, *_ = np.linalg.lstsq(X[rows], y[rows], rcond=None)
    return beta, X


def sse(params, rows=slice(None)):
    beta, X = fit_ols(params, rows)
    resid = y[rows] - X[rows] @ beta
    return float(resid @ resid)

# %% [markdown]
# ## 4. Search for the adstock and saturation parameters
#
# OLS finds the coefficients once the transforms are fixed, so we only need to search the six
# nonlinear parameters: a decay per channel (0 to 0.9) and a half-saturation multiple per channel
# (0.2 to 10 times average spend). A coarse grid gives a starting point, then scipy polishes it.

# %% search
bounds = [(0.0, 0.9)] * 3 + [(0.2, 10.0)] * 3

grid = [(d, k) for d in np.arange(0, 0.91, 0.1) for k in [0.5, 1, 2, 5, 10]]
start = min(([d] * 3 + [k] * 3 for d, k in grid), key=sse)

result = minimize(sse, x0=start, bounds=bounds, method="L-BFGS-B")
best = result.x
beta, X = fit_ols(best)
fitted = X @ beta

params_table = pd.DataFrame({
    "decay": best[:3].round(3),
    "half_sat_x_mean": best[3:].round(2),
    "half_sat_dollars": (best[3:] * mean_spend).round(0),
}, index=CHANNELS)
print(params_table)

r2 = 1 - ((y - fitted) ** 2).sum() / ((y - y.mean()) ** 2).sum()
mape = np.mean(np.abs((y - fitted) / y))
print(f"R-squared: {r2:.4f}   MAPE: {mape:.2%}")
print(pd.Series(beta, index=control_names + CHANNELS).round(1))

# %% [markdown]
# ## 5. Decompose revenue into baseline, promotion and channels

# %% decomposition
contrib = pd.DataFrame(X * beta, columns=control_names + CHANNELS, index=df["week"])
decomp = pd.DataFrame({
    "baseline": contrib[["intercept", "trend", "season_sin", "season_cos"]].sum(axis=1),
    "promotion": contrib["promotion"],
    "search": contrib["search_spend"],
    "social": contrib["social_spend"],
    "video": contrib["video_spend"],
})
share = decomp.sum() / decomp.sum().sum()
print(share.round(3))

fig, ax = plt.subplots(figsize=(10, 4))
ax.stackplot(decomp.index, decomp.T.values, labels=decomp.columns, alpha=0.85)
ax.plot(df["week"], y, color="black", lw=1, label="actual revenue")
ax.set_ylim(15000, None)
ax.legend(loc="upper left", ncol=6, fontsize=8)
ax.set_title("Weekly revenue decomposition (baseline axis starts at 15,000)")
plt.tight_layout()
plt.show()

# %% [markdown]
# ## 6. ROI per channel
#
# Average ROI is modeled contribution divided by spend over the two years. Marginal ROI is the
# return on the next dollar at average spend, read off the slope of the response curve.

# %% roi
decays, half = best[:3], best[3:] * mean_spend
media_beta = beta[-3:]


def weekly_response(s, j):
    """Steady-state weekly contribution if channel j spends s every week."""
    a = s / (1 - decays[j])
    return media_beta[j] * hill(a, half[j])


def marginal_roi(s, j):
    a = s / (1 - decays[j])
    return media_beta[j] * half[j] / (a + half[j]) ** 2 / (1 - decays[j])


roi = pd.DataFrame({
    "spend": spend.sum(axis=0).round(0),
    "contribution": contrib[CHANNELS].sum().values.round(0),
}, index=CHANNELS)
roi["roi"] = (roi["contribution"] / roi["spend"]).round(2)
roi["marginal_roi"] = [round(marginal_roi(mean_spend[j], j), 2) for j in range(3)]
print(roi)

# %% [markdown]
# ## 7. Response curves

# %% curves
fig, ax = plt.subplots(figsize=(8, 4))
for j, name in enumerate(CHANNELS):
    s = np.linspace(0, 3 * mean_spend[j], 100)
    ax.plot(s, weekly_response(s, j), label=name)
    ax.scatter(mean_spend[j], weekly_response(mean_spend[j], j), zorder=3)
ax.set_xlabel("Weekly spend")
ax.set_ylabel("Weekly revenue contribution")
ax.set_title("Response curves (dots = average spend; beyond the highest observed spend is extrapolation)")
ax.legend()
plt.tight_layout()
plt.show()

# %% [markdown]
# ## 8. Reallocate the same budget
#
# Keep the weekly total fixed at today's average and let each channel move at most 30% from its
# average, so the plan stays close to spend levels the model has actually seen.

# %% allocate
total = mean_spend.sum()


def neg_total_response(s):
    return -sum(weekly_response(s[j], j) for j in range(3))


plan = minimize(
    neg_total_response,
    x0=mean_spend,
    bounds=[(0.7 * m, 1.3 * m) for m in mean_spend],
    constraints=[{"type": "eq", "fun": lambda s: s.sum() - total}],
    method="SLSQP",
).x

allocation = pd.DataFrame({
    "current": mean_spend.round(0),
    "proposed": plan.round(0),
    "change": (plan - mean_spend).round(0),
}, index=CHANNELS)
print(allocation)
uplift = -neg_total_response(plan) + neg_total_response(mean_spend)
print(f"Extra weekly revenue for the same budget: {uplift:,.0f}")

# %% [markdown]
# ## 9. Holdout check
#
# Refit everything (parameters and coefficients) on the first 78 weeks, then predict the last 26
# weeks the model has never seen.

# %% holdout
train = slice(0, 78)
test = slice(78, n)
holdout_fit = minimize(sse, x0=start, args=(train,), bounds=bounds, method="L-BFGS-B")
beta_train, X_all = fit_ols(holdout_fit.x, train)
pred = X_all @ beta_train
test_mape = np.mean(np.abs((y[test] - pred[test]) / y[test]))
train_mape = np.mean(np.abs((y[train] - pred[train]) / y[train]))
print(f"Train MAPE: {train_mape:.2%}   Holdout MAPE: {test_mape:.2%}")

fig, ax = plt.subplots(figsize=(10, 3.5))
ax.plot(df["week"], y, color="black", lw=1, label="actual")
ax.plot(df["week"], pred, color="tab:green", label="model trained on first 78 weeks")
ax.axvline(df["week"].iloc[78], ls="--", color="gray")
ax.legend()
plt.tight_layout()
plt.show()

# %% [markdown]
# ## 10. Check against the answer key
#
# This dataset was generated with known effects: each dollar of search, social and video spend
# added 1.4, 0.8 and 0.6 of revenue in the same week, with no carryover and no saturation. A model
# that works should get close to those numbers. On your own data there is no answer key, which is
# why experiments matter.
#
# The half-saturation multiples all finished on the upper bound of 10, which is the model telling
# you it can't see diminishing returns in this data. Widening that bound to 100 shows how much the
# bound itself was shaping the average ROI.

# %% answer_key
truth = pd.Series([1.4, 0.8, 0.6], index=CHANNELS)
wide = minimize(sse, x0=start, bounds=[(0.0, 0.9)] * 3 + [(0.2, 100.0)] * 3, method="L-BFGS-B")
beta_wide, X_wide = fit_ols(wide.x)
roi_wide = (X_wide[:, -3:] * beta_wide[-3:]).sum(axis=0) / spend.sum(axis=0)
print(pd.DataFrame({
    "roi_bound_10": roi["roi"],
    "roi_bound_100": roi_wide.round(2),
    "true_roi": truth,
}))
print("Half-saturation multiples with the wider bound:", wide.x[3:].round(1))