"""PyMC-Marketing quickstart on the Vexpower practice dataset.

Fits a Bayesian marketing mix model with PyMC-Marketing's MMM class
(GeometricAdstock + LogisticSaturation) to mmm-practice-synthetic.csv,
prints the media parameters and ROAS by channel, then asks the budget
optimizer how to split next quarter's weekly budget.

Tested with pymc-marketing 1.3.0 (October 2026) on Python 3.13, CPU only.
Install:  pip install "pymc-marketing==1.3.0"
Run:      python pymc-marketing-quickstart.py [path/to/mmm-practice-synthetic.csv]

The data is synthetic Vexpower teaching data (104 weeks: week, revenue,
search_spend, social_spend, video_spend, promotion), generated from an exactly
linear, noiseless equation (revenue = 20000 + 1.4*search + 0.8*social
+ 0.6*video + 3000*sin(pi*t/26) + 1800*promotion) with no adstock and no
saturation. The true ROAS values are 1.4, 0.8 and 0.6; any fitted carryover
or curvature is not a finding. The sampler settings are
light so the script finishes in a minute or two; use more chains and draws,
and check the convergence warnings, before you act on anything it says.
"""

import sys

import pandas as pd

from pymc_marketing.mmm import MMM, GeometricAdstock, LogisticSaturation

CSV = sys.argv[1] if len(sys.argv) > 1 else "mmm-practice-synthetic.csv"
CHANNELS = ["search_spend", "social_spend", "video_spend"]

df = pd.read_csv(CSV, parse_dates=["week"])
X = df[["week", *CHANNELS, "promotion"]]
y = df["revenue"]

mmm = MMM(
    date_column="week",
    target_column="revenue",
    channel_columns=CHANNELS,
    control_columns=["promotion"],
    adstock=GeometricAdstock(l_max=8),  # normalize=True is the default
    saturation=LogisticSaturation(),  # lam ~ Gamma(3, 1), beta ~ HalfNormal(2)
    yearly_seasonality=2,
)

# Four chains of 1,000 draws runs in about a minute on a laptop. Read the
# convergence warnings PyMC prints (r-hat, effective sample size, tree depth)
# before you trust anything below.
mmm.fit(X, y, chains=4, draws=1000, tune=1000, target_accept=0.95, random_seed=42)

posterior = mmm.idata.posterior
print("\nPosterior means of the media parameters (channel inputs are max-scaled):")
params = pd.DataFrame(
    {
        name: posterior[name].mean(("chain", "draw")).to_series()
        for name in ["adstock_alpha", "saturation_lam", "saturation_beta"]
    }
)
print(params.round(3))

print("\nROAS over the whole 104 weeks:")
roas = mmm.summary.roas(frequency="all_time")
print(roas.round(2).to_string())

# Budget optimisation for the 13 weeks after the data ends. total_budget is
# per period (per week here), so pass the weekly amount you plan to spend.
last_week = df["week"].max()
start = (last_week + pd.Timedelta(weeks=1)).strftime("%Y-%m-%d")
end = (last_week + pd.Timedelta(weeks=13)).strftime("%Y-%m-%d")
weekly_budget = float(df[CHANNELS].tail(13).sum(axis=1).mean())

current = df[CHANNELS].tail(13).mean()

# Keep each channel within 50-150 percent of its current weekly spend, so the
# optimizer can't wander far outside the spend levels the model has seen.
bounds = {ch: (0.5 * current[ch], 1.5 * current[ch]) for ch in CHANNELS}

optimizer = mmm.budget_optimizer(start_date=start, end_date=end)
result = optimizer.allocate_budget(total_budget=weekly_budget, budget_bounds=bounds)

print(f"\nCurrent weekly split (last 13 weeks, total {weekly_budget:,.0f}):")
print(current.round(0).to_string())
print("\nOptimised weekly split:")
print(result.budgets.to_series().round(0).to_string())