Building Statistically Sound Marketing Mix Models

# analysis_plan.py — Document this BEFORE fitting
"""
Analysis Plan: Q4 2025 Media Effectiveness
==========================================
Business Question: What is the incremental ROI of each media channel?
Decision: Inform 2026 budget allocation

Outcome Variable: Weekly national sales ($)
Media Channels:
  - TV (national, geometric adstock, L_max=8)
  - Digital Display (national, geometric adstock, L_max=4)
  - Paid Search (national, geometric adstock, L_max=2)

Confounders (always included):
  - Price index
  - Distribution (% ACV)
  - Holiday indicator
  - Seasonality (Fourier, order 2)

Precision Controls (Bayesian variable selection):
  - Weather (heating degree days)
  - Competitor A spend

Trend: Linear
Sampler: NumPyro, 4 chains, 2000 draws, 1000 tune

Validation: Compare TV ROI posterior to Q3 geo lift test
            (TV ROI from geo test: 1.1-1.5, 90% CI)

Decision Criteria:
  - If TV 80% CI includes geo-test estimate -> validated
  - If TV 80% CI does not overlap -> investigate discrepancy
"""

import pandas as pd
from mmm_framework import (
    ControlVariableConfigBuilder,
    MFFConfigBuilder,
    load_mff,
)

# Configure data structure per analysis plan
mff_config = (
    MFFConfigBuilder()
    .with_kpi_name("Sales")
    .add_national_media("TV", adstock_lmax=8)
    .add_national_media("Digital_Display", adstock_lmax=4)
    .add_national_media("Paid_Search", adstock_lmax=2)
    .add_price_control("Price_Index")          # convenience: allows a negative sign
    .add_distribution_control("Distribution")  # convenience: positive-only
    .add_control(ControlVariableConfigBuilder("Holiday").national().build())
    .build()
)

# Load and validate
panel = load_mff(pd.read_csv("data/weekly_mff.csv"), mff_config)

# Data quality checks (plain pandas on the panel's aligned frames)
print(f"Date range: {panel.coords.periods.min()} to {panel.coords.periods.max()}")
print(f"Observations: {panel.n_obs}")
print(f"Missing media values:\n{panel.X_media.isna().sum()}")
print(f"Missing control values:\n{panel.X_controls.isna().sum()}")
print(f"\nZero-spend share by channel:")
print((panel.X_media == 0).mean())  # flag channels with little variation
print(panel.summary())              # overall panel summary

from mmm_framework import (
    ModelConfigBuilder,
    TrendConfig,
    TrendType,
    BayesianMMM,
)

# Model inference configuration
model_config = (
    ModelConfigBuilder()
    .bayesian_numpyro()       # JAX-based sampler for speed
    .with_chains(4)           # 4 chains for convergence diagnostics
    .with_draws(2000)         # 2000 posterior draws per chain
    .with_tune(1000)          # 1000 warmup iterations
    .with_target_accept(0.9)  # Target acceptance rate
    .build()
)

# Trend configuration (from analysis plan: linear)
trend_config = TrendConfig(
    type=TrendType.LINEAR,
    growth_prior_sigma=0.1
)

# Build the model
mmm = BayesianMMM(panel, model_config, trend_config)

# Verify model structure matches analysis plan
print("Model parameters:")
for var in mmm.model.free_RVs:
    print(f"  {var.name}")

from mmm_framework import (
    AdstockConfigBuilder,
    MediaChannelConfigBuilder,
    PriorConfigBuilder,
)

# Adstock decay priors: encode belief about carryover duration.
# These apply on the parametric-adstock path (see the note below).
# TV: expect longer carryover (allow data to inform)
tv_adstock = (
    AdstockConfigBuilder()
    .geometric()
    .with_max_lag(8)
    .with_alpha_prior(
        PriorConfigBuilder().beta(alpha=3, beta=1.5).build()
        # Beta(3, 1.5) -> mode at 0.8, most mass above 0.5
    )
    .build()
)

# Digital: expect shorter carryover
digital_adstock = (
    AdstockConfigBuilder()
    .geometric()
    .with_max_lag(4)
    .with_alpha_prior(
        PriorConfigBuilder().beta(alpha=2, beta=3).build()
        # Beta(2, 3) -> mode at ~0.33, most mass below 0.6
    )
    .build()
)

# Attach to the channel definitions used in the MFF config
tv_channel = (
    MediaChannelConfigBuilder("TV")
    .national()
    .with_adstock(tv_adstock)
    .with_logistic_saturation()
    .build()
)

# Document: "Multi-week TV carryover is a common industry rule of
# thumb -- treat it as a prior starting point, not evidence.
# Digital prior based on platform-reported attribution windows
# of 1-3 weeks."

# Sample from priors only (no likelihood evaluation yet)
prior = mmm.get_prior(samples=500, random_seed=42)

# y_obs is modeled on a standardized scale; y_obs_scaled is the
# deterministic back-transform to the original KPI scale
y_prior = prior.prior["y_obs_scaled"].values.flatten()

print(f"Prior predictive y range: [{y_prior.min():.0f}, {y_prior.max():.0f}]")
print(f"Actual y range: [{panel.y.min():.0f}, {panel.y.max():.0f}]")
print(f"Prior predictive y mean: {y_prior.mean():.0f}")
print(f"Actual y mean: {panel.y.mean():.0f}")

# Check: do priors produce data in the right ballpark?
# They should cover the observed range with room to spare,
# but not predict absurdities (negative sales, 10x actual)

# Fit with a fixed random seed for reproducibility
results = mmm.fit(random_seed=42)

n_chains = results.trace.posterior.sizes["chain"]
n_draws = results.trace.posterior.sizes["draw"]

print(f"Sampling completed.")
print(f"  Chains: {n_chains}")
print(f"  Draws per chain: {n_draws}")
print(f"  Total posterior samples: {n_chains * n_draws}")

# Convergence diagnostics (results.diagnostics is filled at fit time
# with divergences, rhat_max, and ess_bulk_min)
print("=== MCMC Diagnostics ===")
diag = results.diagnostics

# 1. Divergences: should be 0
print(f"Divergences: {diag['divergences']}")
if diag['divergences'] > 0:
    print("  ACTION: Reparameterize or increase target_accept")

# 2. R-hat: should be < 1.01 for all parameters
print(f"R-hat max: {diag['rhat_max']:.4f}")
if diag['rhat_max'] > 1.01:
    print("  ACTION: Run longer chains or investigate multimodality")

# 3. ESS: should be > 400 for all parameters
# (tail ESS is not pre-computed; pull it from the trace with ArviZ)
import arviz as az
# arviz 1.x returns a DataTree (no .to_array()); dataset_extremum reduces it to a float
from mmm_framework.utils.arviz_compat import dataset_extremum
ess_tail_min = dataset_extremum(az.ess(results.trace, method="tail"), "min")
print(f"ESS bulk min: {diag['ess_bulk_min']:.0f}")
print(f"ESS tail min: {ess_tail_min:.0f}")
if diag['ess_bulk_min'] < 400:
    print("  ACTION: Run more draws")

# 4. Summary of key parameters
summary = results.summary()
print("\n=== Parameter Summary ===")
# arviz 1.x exposes an 89% equal-tailed interval as eti89_lb / eti89_ub
print(summary[["mean", "sd", "eti89_lb", "eti89_ub", "r_hat"]].to_string())

# Posterior predictive check via the model's predict() method
import numpy as np

pred = mmm.predict(hdi_prob=0.90, random_seed=42)

# Compare predicted vs actual (original KPI scale)
y_pred_mean = pred.y_pred_mean
y_actual = panel.y.values

# Calibration: what fraction of observations fall within the 90% HDI?
coverage = np.mean(
    (y_actual >= pred.y_pred_hdi_low) & (y_actual <= pred.y_pred_hdi_high)
)

print(f"90% interval coverage: {coverage:.1%}")
print(f"  (Target: ~90%. If much lower, model is overconfident)")
print(f"  (If much higher, model is underconfident)")

# MAPE
mape = np.mean(np.abs(y_pred_mean - y_actual) / y_actual)
print(f"MAPE: {mape:.1%}")

# Out-of-sample backtest: genuine out-of-time forecasts
from mmm_framework.validation import BacktestConfig, run_backtest

config = BacktestConfig(
    min_train_size=104,  # two seasonal cycles before the first cutoff
    horizon=13,          # forecast a quarter past each cutoff
    step=13,             # non-overlapping windows
)
result = run_backtest(mmm, config)

print(result.fits)          # per-refit convergence -- read this first
print(result.summary())     # MAPE/MASE vs naive + seasonal-naive baselines
print(result.coverage())    # nominal vs empirical interval coverage

# Gates worth pre-registering:
#   - MASE < 1 (beats "copy last year" out of time)
#   - empirical coverage near nominal (intervals honest on unseen data)

# Sensitivity analysis: vary key assumptions
sensitivity_results = {}

def tv_roas(model, panel):
    """Posterior TV ROAS: counterfactual contribution / spend."""
    contrib = model.compute_counterfactual_contributions(
        channels=["TV"], hdi_prob=0.9, random_seed=42
    )
    spend = panel.X_media["TV"].sum()
    return (
        contrib.total_contributions["TV"] / spend,
        contrib.contribution_hdi_low["TV"] / spend,
        contrib.contribution_hdi_high["TV"] / spend,
    )

# 1. Vary TV adstock L_max
for lmax in [4, 6, 8, 12]:
    config = make_config(tv_lmax=lmax)  # your helper: rebuild configs with this l_max
    model = BayesianMMM(panel, config, trend_config)
    model.fit(random_seed=42)
    sensitivity_results[f"tv_lmax_{lmax}"] = tv_roas(model, panel)

# 2. Vary prior strength
for sigma in [0.5, 1.0, 2.0]:
    config = make_config(beta_prior_sigma=sigma)
    model = BayesianMMM(panel, config, trend_config)
    model.fit(random_seed=42)
    sensitivity_results[f"prior_sigma_{sigma}"] = tv_roas(model, panel)

# Report: how much do ROAS estimates change?
print("\n=== Sensitivity Analysis ===")
for name, (roas, low, high) in sensitivity_results.items():
    print(f"  {name}: TV ROAS = {roas:.2f} ({low:.2f}-{high:.2f})")

# Channel contributions and ROAS from counterfactual analysis:
# contribution = prediction(actual spend) - prediction(channel zeroed out)
contrib = mmm.compute_counterfactual_contributions(hdi_prob=0.9, random_seed=42)

print("=== Channel Results ===")
for channel in panel.coords.channels:
    spend = panel.X_media[channel].sum()
    roas = contrib.total_contributions[channel] / spend
    roas_low = contrib.contribution_hdi_low[channel] / spend
    roas_high = contrib.contribution_hdi_high[channel] / spend
    print(f"\n{channel}:")
    print(f"  ROAS: {roas:.2f} (90% HDI: {roas_low:.2f}-{roas_high:.2f})")
    print(f"  Share of media effect: {contrib.contribution_pct[channel]:.1f}%")

# Generate the full report
from mmm_framework.reporting import MMMReportGenerator, ReportConfig

report = MMMReportGenerator(
    model=mmm,
    panel=panel,
    results=results,
    config=ReportConfig(
        title="Q4 2025 Media Effectiveness Analysis",
        client="Brand Name",
        analysis_period="Jan 2024 - Dec 2025",
    ),
)

report.to_html("q4_2025_mmm_report.html")