Getting Started

import mmm_framework
print(f"mmm-framework version: {mmm_framework.__version__}")

# Check available components
from mmm_framework import (
    MFFConfigBuilder,
    ModelConfigBuilder,
    BayesianMMM,
    load_mff,
)
print("✓ All core components imported successfully")

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

# 1. Load a bundled example — 104 weeks of national weekly data, ready to model.
panel = load_example("national")
print(panel.summary())

# 2. Configure inference (fast JAX/NumPyro sampler) and a linear trend.
model_config = (
    ModelConfigBuilder()
    .bayesian_numpyro()      # ~3x faster than PyMC at equal draws
    .with_chains(4)
    .with_draws(500)         # small + fast for a first run
    .with_tune(500)
    .build()
)
trend_config = TrendConfig(type=TrendType.LINEAR)

# 3. Fit. On a laptop this national model takes roughly 15-25 seconds.
mmm = BayesianMMM(panel, model_config, trend_config)
results = mmm.fit(random_seed=42)
print("max R-hat:", round(results.diagnostics["rhat_max"], 3))   # ~1.0 = converged

# 4. The headline: each channel's return on ad spend (contribution / spend).
decomp = mmm.compute_component_decomposition()
roi = (decomp.media_by_channel.sum() / panel.X_media.sum()).sort_values(ascending=False)
print("\nEstimated ROI by channel:")
print(roi.round(2))

# 5. This example ships a SEALED answer key — grade the estimate against truth.
truth = load_example_answer_key("national")["true_roas"]
print("\nchannel   estimated   true")
for ch in roi.index:
    print(f"  {ch:<8} {roi[ch]:>8.2f}   {truth[ch]:>5.2f}")

import pandas as pd
import numpy as np

# Example: Create synthetic MFF data
np.random.seed(42)
n_weeks = 104  # 2 years of weekly data

# Generate dates
dates = pd.date_range("2023-01-01", periods=n_weeks, freq="W")

# Build MFF records
records = []

# KPI: Sales (weekly)
for i, date in enumerate(dates):
    base_sales = 1000 + 50 * np.sin(2 * np.pi * i / 52)  # Seasonality
    noise = np.random.normal(0, 50)
    records.append({
        "Period": date,
        "VariableName": "Sales",
        "VariableValue": base_sales + noise
    })

# Media: TV spend (weekly, with some weeks at zero)
for i, date in enumerate(dates):
    spend = np.random.exponential(500) if np.random.random() > 0.2 else 0
    records.append({
        "Period": date,
        "VariableName": "TV",
        "VariableValue": spend
    })

# Media: Digital spend (weekly)
for i, date in enumerate(dates):
    spend = np.random.exponential(300)
    records.append({
        "Period": date,
        "VariableName": "Digital",
        "VariableValue": spend
    })

# Control: Price Index
for i, date in enumerate(dates):
    price = 100 + np.random.normal(0, 5)
    records.append({
        "Period": date,
        "VariableName": "Price",
        "VariableValue": price
    })

# Create DataFrame
mff_data = pd.DataFrame(records)
print(f"MFF shape: {mff_data.shape}")
print(mff_data.head(10))

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

# Step 2a: Configure the data structure
mff_config = (
    MFFConfigBuilder()
    .with_kpi_name("Sales")                    # Target variable
    .add_national_media("TV", adstock_lmax=8)  # TV with 8-week carryover
    .add_national_media("Digital", adstock_lmax=4)  # Digital with 4-week carryover
    .add_price_control()                       # Price as control variable
    .build()
)

print(f"Media channels: {mff_config.media_names}")
print(f"Control variables: {mff_config.control_names}")

# Step 2b: Load and validate data
panel = load_mff(mff_data, mff_config)

print(f"\nPanel dataset:")
print(f"  Observations: {panel.n_obs}")
print(f"  Channels: {panel.n_channels}")
print(f"  Controls: {panel.n_controls}")

# Step 2c: Configure model inference
model_config = (
    ModelConfigBuilder()
    .bayesian_numpyro()       # JAX sampler (measured ~3x faster than PyMC at equal draws on PyMC 6)
    .with_chains(4)           # 4 parallel chains for convergence diagnostics
    .with_draws(2000)         # 2000 posterior draws
    .with_tune(1000)          # 1000 warmup iterations
    .with_target_accept(0.9)  # Target acceptance rate
    .build()
)

# Step 2d: Configure trend component
trend_config = TrendConfig(
    type=TrendType.LINEAR,
    growth_prior_sigma=0.1
)

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

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

# Prior predictive check (ALWAYS do this before fitting)
print("\n=== Prior Predictive Check ===")
prior = mmm.sample_prior_predictive(samples=200)
y_prior = prior.prior_predictive["y_obs"].values.flatten()

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

# Fit the model
print("\n=== Fitting Model ===")
results = mmm.fit(random_seed=42)

# Convergence diagnostics (CRITICAL - check these!)
print("\n=== Diagnostics ===")
print(f"Divergences: {results.diagnostics['divergences']}")
print(f"R-hat max: {results.diagnostics['rhat_max']:.4f}")
print(f"ESS bulk min: {results.diagnostics['ess_bulk_min']:.0f}")

# Check for issues
if results.diagnostics['divergences'] > 0:
    print("⚠️  Divergences detected - consider reparameterization")
if results.diagnostics['rhat_max'] > 1.01:
    print("⚠️  R-hat > 1.01 - chains may not have converged")
if results.diagnostics['ess_bulk_min'] < 400:
    print("⚠️  Low ESS - consider more draws")

# Posterior summary with uncertainty
print("\n=== Posterior Summary ===")
summary = results.summary(["beta_TV", "beta_Digital", "sigma"])
# ArviZ 1.x reports an 89% equal-tailed interval as eti89_lb / eti89_ub
# (the older hdi_3% / hdi_97% column names were removed).
print(summary[["mean", "sd", "eti89_lb", "eti89_ub", "r_hat"]])

# Channel contributions with uncertainty
print("\n=== Channel Contributions ===")
if results.channel_contributions is not None:
    contrib = results.channel_contributions.sum()
    print(f"Total TV contribution: {contrib['TV']:.0f}")
    print(f"Total Digital contribution: {contrib['Digital']:.0f}")

# Access the full posterior trace (an ArviZ DataTree under arviz 1.x) for plots
import arviz as az
from mmm_framework.utils.arviz_compat import plot_posterior

# Posterior distributions — az.plot_posterior was removed in arviz 1.x, so route
# through the framework's version-robust shim (it calls arviz_plots.plot_dist)
plot_posterior(results.trace, var_names=["beta_TV", "beta_Digital"])

# Trace plots for convergence
az.plot_trace(results.trace, var_names=["beta_TV", "beta_Digital"])

# Forest plot comparing channels
az.plot_forest(results.trace, var_names=["beta_TV", "beta_Digital"])

from mmm_framework.reporting import MMMReportGenerator, ReportConfig, ReportBuilder

# Option 1: Quick report from fitted model
report = MMMReportGenerator(
    model=mmm,
    panel=panel,
    results=results,
    config=ReportConfig(
        title="Marketing Mix Model Analysis",
        client="Acme Consumer Products",
        analysis_period="Jan 2023 - Dec 2024",
    ),
)

# Save to HTML
report.to_html("mmm_report_q4_2024.html")

# Option 2: Fluent builder pattern for customization
report = (
    ReportBuilder()
    .with_model(mmm, panel=panel, results=results)
    .with_title("Q4 Marketing Analysis")
    .with_client("Acme Corp")
    .with_credible_interval(0.9)  # 90% credible intervals
    .enable_all_sections()
    .disable_section("diagnostics")  # Hide technical details
    .build()
)

report.to_html("executive_summary.html")

# MFF structure example
"""
Period      | Geography | VariableName | VariableValue
2024-01-01  | National  | Sales        | 15000
2024-01-01  | National  | TV_Spend     | 50000
2024-01-01  | East      | Sales        | 8000
2024-01-01  | West      | Sales        | 7000
2024-01-08  | National  | Sales        | 16500
...
"""

# The framework handles dimension alignment automatically
mff_config = (
    MFFConfigBuilder()
    .with_kpi_name("Sales")
    # National media disaggregated to geo by population share
    .add_national_media("TV", adstock_lmax=8)
    # Geo-level media stays at its native granularity
    .add_media_builder(
        MediaChannelConfigBuilder("Local_Radio").by_geo().with_geometric_adstock(6)
    )
    .build()
)

from mmm_framework import (
    PriorConfigBuilder,
    AdstockConfigBuilder,
    SaturationConfigBuilder,
    MediaChannelConfigBuilder,
    HierarchicalConfigBuilder,
    SeasonalityConfigBuilder,
)

# Build custom priors
decay_prior = (
    PriorConfigBuilder()
    .beta(alpha=2, beta=2)  # Centered at 0.5
    .build()
)

# Build adstock configuration
adstock = (
    AdstockConfigBuilder()
    .geometric()
    .with_max_lag(8)
    .with_alpha_prior(decay_prior)
    .build()
)

# Build complete media channel config
tv_channel = (
    MediaChannelConfigBuilder("TV")  # name is a required constructor argument
    .national()
    .with_adstock(adstock)
    .with_hill_saturation()
    .with_positive_prior()  # Constrain the coefficient to be positive (HalfNormal)
    .build()
)

# Build hierarchical structure for geo models
hierarchical = (
    HierarchicalConfigBuilder()
    .enabled()
    .pool_across_geo()
    .use_non_centered()  # Better for sparse geos
    .with_non_centered_threshold(20)
    .build()
)