"""Google Meridian quickstart on the Vexpower practice dataset.

Fits a national-level Meridian model to mmm-practice-synthetic.csv with ROI
priors (the library default for two channels, a tighter "experiment" prior for
search), checks convergence, prints ROI by channel and runs the default
fixed-budget optimisation.

Tested with google-meridian 2.1.0 (October 2026) on Python 3.13, CPU only.
Install:  pip install "google-meridian==2.1.0"
          (Linux with an NVIDIA GPU: pip install "google-meridian[and-cuda]==2.1.0")
Run:      python meridian-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). It is national, so this
is a single-geo model, which throws away Meridian's main strength (geo-level
hierarchy). The file was 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, so the true ROIs are
1.4, 0.8 and 0.6 and any fitted decay or half-saturation is not a finding.
The sampler settings are light so it finishes on a laptop CPU;
Google's own demo uses 10 chains and 1,000 kept draws on a GPU.
"""

import sys

import numpy as np
import pandas as pd

import tensorflow_probability.substrates.jax as tfp_jax  # JAX is the default backend
from meridian import constants
from meridian.analysis import analyzer, optimizer, visualizer
from meridian.data import data_frame_input_data_builder
from meridian.model import model, prior_distribution, spec

CSV = sys.argv[1] if len(sys.argv) > 1 else "mmm-practice-synthetic.csv"
CHANNELS = ["search", "social", "video"]

df = pd.read_csv(CSV)
df = df.rename(columns={"week": "time"})  # Meridian's default time column name

# 1. Data. kpi_type="revenue" means the KPI is already money, so ROI is
#    revenue per unit of spend. We pass spend as the media metric too, since
#    the practice file has no impressions.
builder = data_frame_input_data_builder.DataFrameInputDataBuilder(
    kpi_type="revenue",
    default_kpi_column="revenue",
)
builder = (
    builder.with_kpi(df)
    .with_controls(df, control_cols=["promotion"])
    .with_media(
        df,
        media_cols=[f"{c}_spend" for c in CHANNELS],
        media_spend_cols=[f"{c}_spend" for c in CHANNELS],
        media_channels=CHANNELS,
    )
)
data = builder.build()

# 2. Priors. Meridian's default ROI prior is LogNormal(0.2, 0.9) for every
#    channel (median ROI about 1.22). Here search gets an informative prior,
#    as if an experiment had found ROI between 1.2 and 2.0 with 90 percent
#    probability. Order follows CHANNELS.
search_prior = prior_distribution.lognormal_dist_from_range(1.2, 2.0, mass_percent=0.9)
mu = np.array([float(search_prior.loc), 0.2, 0.2])
sigma = np.array([float(search_prior.scale), 0.9, 0.9])
prior = prior_distribution.PriorDistribution(
    roi_m=tfp_jax.distributions.LogNormal(mu, sigma, name=constants.ROI_M)
)
model_spec = spec.ModelSpec(prior=prior, max_lag=8)

mmm = model.Meridian(input_data=data, model_spec=model_spec)

# 3. Sample. Light settings for a CPU; raise n_chains / n_keep for real work.
mmm.sample_prior(200)
mmm.sample_posterior(n_chains=4, n_adapt=500, n_burnin=500, n_keep=500, seed=0)

# 4. Convergence: r-hat close to 1.0 is what you want (Meridian's docs treat
#    values under 1.2 as approximate convergence).
print("\nR-hat summary by parameter:")
mmm_analyzer = analyzer.Analyzer(
    model_context=mmm.model_context, inference_data=mmm.inference_data
)
print(mmm_analyzer.rhat_summary().round(3).to_string())

# 5. Results: ROI, contribution and spend by channel.
media_summary = visualizer.MediaSummary(mmm)
print("\nMedia summary:")
print(media_summary.summary_table().to_string())

# 6. Default fixed-budget optimisation over the historical period, with each
#    channel allowed to move 30 percent either way (the library default).
opt = optimizer.BudgetOptimizer(mmm)
results = opt.optimize()
print("\nHistorical vs optimised allocation (posterior means):")
print(results.nonoptimized_data[["spend", "roi"]].sel(metric="mean").to_dataframe().round(2).to_string())
print(results.optimized_data[["spend", "roi"]].sel(metric="mean").to_dataframe().round(2).to_string())