"""Turn one lift test into calibration inputs for Robyn, PyMC-Marketing and Meridian.

From the Vexpower MMM library: vexpower.com/mmm/calibrating-mmm-with-experiments

Edit the EXPERIMENT block, then run:
    python calibration_inputs.py [path/to/weekly_data.csv]

With no argument it looks for mmm-practice-synthetic.csv (download it from
vexpower.com/downloads/mmm/mmm-practice-synthetic.csv) to read the channel's
average weekly spend. It prints:
  1. the experiment's incremental return and standard error,
  2. an R snippet for Robyn's calibration_input,
  3. a pandas DataFrame row for PyMC-Marketing's add_lift_test_measurements,
  4. LogNormal ROI prior parameters and a Python snippet for Meridian,
  5. a precision-weighted comparison with your current MMM estimate.

Only numpy and pandas are needed. Nothing here calls the libraries themselves;
check each library's current docs before relying on the argument names.
"""
import math
import sys
from pathlib import Path

import numpy as np
import pandas as pd

# ---------------------------------------------------------------- inputs ---
EXPERIMENT = {
    "channel": "social_spend",          # column name in your MMM data
    "start": "2025-09-01",              # first day of the test
    "end": "2025-09-28",                # last day counted (include any cool-down)
    "incremental_revenue": 3900.0,      # lift in the KPI the MMM models (here revenue)
    "incremental_revenue_se": 1250.0,   # its standard error
    "spend_in_test": 2800.0,            # spend that caused the lift (only the tested cells)
    "kpi": "revenue",                   # must match the MMM's dependent variable
    "counterfactual": "zero",           # "zero" (ads off in test cells) or "reduced" (budget cut, not off)
    "spend_change": "removed",          # "removed" (holdout / go-dark) or "added" (heavy-up)
    "weeks": 4,
}
# What the MMM currently says about this channel's ROI, with a standard error.
MMM_ROI, MMM_ROI_SE = 2.1, 0.7
# Extra uncertainty for transferring the result (1.0 = identical estimands).
TRANSFER_FACTOR = 1.3
# ---------------------------------------------------------------------------


def load_weekly_spend(path: Path, channel: str) -> float:
    df = pd.read_csv(path)
    if channel not in df.columns:
        raise SystemExit(f"{channel!r} is not a column in {path}; columns are {list(df.columns)}")
    return float(df[channel].mean())


def lognormal_from_mean_sd(mean: float, sd: float) -> tuple[float, float]:
    """mu and sigma of a LogNormal with the given mean and standard deviation."""
    sigma2 = math.log(1 + (sd / mean) ** 2)
    return math.log(mean) - sigma2 / 2, math.sqrt(sigma2)


def main() -> None:
    exp = EXPERIMENT
    if exp["incremental_revenue"] <= 0:
        raise SystemExit("A non-positive lift can't become a LogNormal ROI prior; record it, don't force it.")
    roi = exp["incremental_revenue"] / exp["spend_in_test"]
    roi_se = exp["incremental_revenue_se"] / exp["spend_in_test"]
    z = exp["incremental_revenue"] / exp["incremental_revenue_se"]
    p_two_sided = math.erfc(abs(z) / math.sqrt(2))

    print("1. Experiment")
    print(f"   iROAS {roi:.2f} (SE {roi_se:.2f}), z = {z:.2f}, two-sided p = {p_two_sided:.3f}")
    if exp["counterfactual"] != "zero":
        print("   WARNING: a reduced-budget test measures a marginal return; the MMM's ROI is against zero spend.")

    print("\n2. Robyn: calibration_input (R)")
    print(
        "calibration_input <- data.frame(\n"
        f'  channel = "{exp["channel"]}",\n'
        f'  liftStartDate = as.Date("{exp["start"]}"),\n'
        f'  liftEndDate = as.Date("{exp["end"]}"),\n'
        f'  liftAbs = {exp["incremental_revenue"]:.0f},\n'
        f'  spend = {exp["spend_in_test"]:.0f},  # must match spend in dt_input for these dates\n'
        f'  confidence = {1 - p_two_sided:.3f},  # Robyn demo: 1 - p-value for a frequentist test\n'
        f'  metric = "{exp["kpi"]}",\n'
        '  calibration_scope = "immediate"  # Robyn recommends "immediate" for experiments\n'
        ")"
    )

    data = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(__file__).with_name("mmm-practice-synthetic.csv")
    if not data.exists():
        data = Path(__file__).resolve().parents[2] / "public/downloads/mmm/mmm-practice-synthetic.csv"
    x = load_weekly_spend(data, exp["channel"]) if data.exists() else float("nan")
    weekly_delta_x = exp["spend_in_test"] / exp["weeks"]
    sign = -1 if exp["spend_change"] == "removed" else 1
    lift_df = pd.DataFrame([{
        "channel": exp["channel"],
        "x": x,
        "delta_x": sign * weekly_delta_x,
        "delta_y": sign * exp["incremental_revenue"] / exp["weeks"],
        "sigma": exp["incremental_revenue_se"] / exp["weeks"],
    }])
    print("\n3. PyMC-Marketing: rows for mmm.add_lift_test_measurements(df_lift_test)")
    print(f"   (average weekly {exp['channel']} in {data.name if data.exists() else 'data'}: {x:,.0f})")
    print(lift_df.to_string(index=False))
    print("   x is the weekly spend level the test started from, delta_x the weekly change in spend,")
    print("   delta_y the weekly change in the KPI and sigma its SE, in original units. Removing the")
    print("   tested spend from the usual level is read as a move down the saturation curve; that is")
    print("   an approximation when the test only covered some users or regions. Check the current docs.")

    sd_prior = roi_se * TRANSFER_FACTOR
    mu, sigma = lognormal_from_mean_sd(roi, sd_prior)
    draws = np.random.default_rng(1).lognormal(mu, sigma, 200_000)
    print("\n4. Meridian: LogNormal ROI prior")
    print(f"   mean {roi:.2f}, sd {sd_prior:.2f} (experiment SE x transfer factor {TRANSFER_FACTOR})")
    print(f"   LogNormal(mu={mu:.3f}, sigma={sigma:.3f}); 90% of mass between "
          f"{np.quantile(draws, 0.05):.2f} and {np.quantile(draws, 0.95):.2f}")
    print(
        "   from meridian.model import prior_distribution\n"
        f"   roi_prior = prior_distribution.lognormal_dist_from_mean_std({roi:.2f}, {sd_prior:.2f})\n"
        "   prior = prior_distribution.PriorDistribution(roi_m=roi_prior)  # one value per channel in practice"
    )

    w_e, w_m = 1 / sd_prior**2, 1 / MMM_ROI_SE**2
    comb = (w_e * roi + w_m * MMM_ROI) / (w_e + w_m)
    comb_se = math.sqrt(1 / (w_e + w_m))
    gap = abs(roi - MMM_ROI) / math.sqrt(sd_prior**2 + MMM_ROI_SE**2)
    print("\n5. Precision-weighted comparison with the current MMM")
    print(f"   MMM {MMM_ROI:.2f} +/- {1.96 * MMM_ROI_SE:.2f}; experiment {roi:.2f} +/- {1.96 * sd_prior:.2f}")
    print(f"   combined {comb:.2f} +/- {1.96 * comb_se:.2f} (weight on experiment {w_e / (w_e + w_m):.0%}); gap = {gap:.1f} SE")
    if gap > 2:
        print("   The estimates disagree by more than 2 SE: check they measure the same thing before combining.")


if __name__ == "__main__":
    main()