"""Jin et al. (2017) transforms, in plain numpy.

Implements, from "Bayesian Methods for Media Mix Modeling with Carryover and
Shape Effects" (Jin, Wang, Sun, Chan & Koehler, Google, 2017):

  * eq. (1)-(3): normalized adstock over L weeks with geometric weights
    alpha**l or delayed weights alpha**((l - theta)**2)
  * eq. (4)-(5): the Hill function beta / (1 + (x / K)**-S)
  * both orders: adstock then Hill (the paper's choice) and Hill then adstock
  * eq. (10)-(11): ROAS and mROAS from counterfactual predictions, summing the
    change period plus the L - 1 weeks after it

Then it fits a small illustrative model to the Vexpower practice CSV
(mmm-practice-synthetic.csv, invented data) by grid search plus least squares.
That fit is NOT the paper's Bayesian estimation: it is here so you can see the
transforms and the ROAS definitions working on a real-shaped dataset. For the
Bayesian version use Meridian or PyMC-Marketing.

Be aware that the practice CSV was generated as an exactly linear model with no
noise, no adstock and no saturation (revenue = 20000 + 1.4*search + 0.8*social
+ 0.6*video + seasonality + promotion). So the fit should find carryover near
zero and push the half-saturation point K to the top of its grid (a nearly
straight line), and mROAS should land close to the true 1.4 / 0.8 / 0.6. Any
fitted decay or curvature here is not a finding about the data.

Usage:
  python jin-2017-transforms.py                      # looks for the CSV next to this file
  python jin-2017-transforms.py path/to/practice.csv
Requires numpy and pandas.
"""
import itertools
import sys
from pathlib import Path

import numpy as np
import pandas as pd


# ---------- the paper's transforms ----------

def lag_weights(alpha, L, theta=None):
    """Normalized weights for lags 0..L-1. theta=None gives geometric adstock."""
    lags = np.arange(L)
    w = alpha ** lags if theta is None else alpha ** ((lags - theta) ** 2)
    return w / w.sum()


def adstock(x, alpha, L=13, theta=None):
    """Eq. (1): weighted average of this week's and the previous L-1 weeks' spend."""
    x = np.asarray(x, dtype=float)
    w = lag_weights(alpha, L, theta)
    out = np.zeros_like(x)
    for lag, wl in enumerate(w):
        out[lag:] += wl * x[: len(x) - lag]
    return out


def hill(x, K, S):
    """Eq. (4): 1 / (1 + (x/K)^-S), equal to 0.5 at x = K."""
    x = np.asarray(x, dtype=float)
    with np.errstate(divide="ignore"):
        return np.where(x > 0, 1.0 / (1.0 + (np.maximum(x, 1e-12) / K) ** (-S)), 0.0)


def media_effect(x, alpha, K, S, beta=1.0, L=13, theta=None, order="adstock_first"):
    if order == "adstock_first":
        return beta * hill(adstock(x, alpha, L, theta), K, S)
    return beta * adstock(hill(x, K, S), alpha, L, theta)


# ---------- the widget's comparison, as a table ----------

def order_comparison(alpha=0.8, theta=3, L=13, K=0.3, S=2.0):
    flights = {
        "4-week burst": np.r_[np.ones(4), np.zeros(L + 4)],
        "12 weeks always-on at 50%": np.r_[np.full(12, 0.5), np.zeros(L + 4)],
        "pulsing 1 on / 1 off": np.r_[np.tile([1.0, 0.0], 6), np.zeros(L + 4)],
        "one heavy week": np.r_[1.0, np.zeros(L + 4)],
    }
    rows = []
    for name, x in flights.items():
        a = media_effect(x, alpha, K, S, L=L, theta=theta, order="adstock_first").sum()
        h = media_effect(x, alpha, K, S, L=L, theta=theta, order="hill_first").sum()
        rows.append({"flight": name, "adstock_then_hill": round(a, 3), "hill_then_adstock": round(h, 3), "difference": f"{a / h - 1:+.1%}"})
    return pd.DataFrame(rows)


# ---------- ROAS and mROAS, eq. (10) and (11) ----------

def roas_mroas(predict, spend, channel, t0, t1, L=13):
    """predict(spend_dict) -> predicted revenue per week. Change period is t0..t1 inclusive."""
    end = min(len(spend[channel]), t1 + L)  # change period plus L-1 weeks after
    base = predict(spend)
    off = {k: v.copy() for k, v in spend.items()}
    off[channel][t0 : t1 + 1] = 0.0
    up = {k: v.copy() for k, v in spend.items()}
    up[channel][t0 : t1 + 1] *= 1.01
    period_spend = spend[channel][t0 : t1 + 1].sum()
    roas = (base[t0:end] - predict(off)[t0:end]).sum() / period_spend
    mroas = (predict(up)[t0:end] - base[t0:end]).sum() / (0.01 * period_spend)
    return roas, mroas


# ---------- illustrative fit on the practice CSV ----------

def fit_practice(csv_path, L=13):
    df = pd.read_csv(csv_path)
    channels = ["search_spend", "social_spend", "video_spend"]
    y = df["revenue"].to_numpy(float)
    n = len(df)
    trend = np.arange(n) / n
    season = np.c_[np.sin(2 * np.pi * np.arange(n) / 52), np.cos(2 * np.pi * np.arange(n) / 52)]
    controls = np.c_[np.ones(n), trend, season, df["promotion"].to_numpy(float)]
    scale = {c: df[c].max() for c in channels}  # spend scaled to 0-1 as in the paper
    spend = {c: df[c].to_numpy(float) for c in channels}

    grid = list(itertools.product([0.0, 0.1, 0.3, 0.5, 0.7, 0.85], [0, 1, 2, 3], [0.2, 0.4, 0.7, 1.0, 1.5, 3.0, 10.0]))  # alpha, theta, K (S fixed at 1: reach curve)
    params = {c: (0.5, 0, 0.5) for c in channels}

    def design(p, sp):
        cols = [media_effect(sp[c] / scale[c], a, K, 1.0, L=L, theta=th) for c, (a, th, K) in p.items()]
        return np.c_[controls, np.column_stack(cols)]

    def sse(p):
        X = design(p, spend)
        coef, *_ = np.linalg.lstsq(X, y, rcond=None)
        if (coef[-len(channels):] < 0).any():
            return np.inf, coef
        return ((y - X @ coef) ** 2).sum(), coef

    for _ in range(3):  # coordinate search, one channel at a time
        for c in channels:
            best = min(grid, key=lambda g: sse({**params, c: g})[0])
            params[c] = best
    err, coef = sse(params)
    predict = lambda sp: design(params, sp) @ coef

    yhat = predict(spend)
    r2 = 1 - ((y - yhat) ** 2).sum() / ((y - y.mean()) ** 2).sum()
    # Paper's setup: the most recent 12 weeks are the post-change period and the year before is the change period.
    t1 = n - 13
    t0 = t1 - 51
    rows = []
    for c in channels:
        a, th, K = params[c]
        roas, mroas = roas_mroas(predict, spend, c, t0, t1, L)
        rows.append({"channel": c, "alpha": a, "theta": th, "K_scaled": K, "beta": round(coef[controls.shape[1] + channels.index(c)], 1), "ROAS": round(roas, 2), "mROAS": round(mroas, 2)})
    return pd.DataFrame(rows), r2, (t0, t1)


if __name__ == "__main__":
    print("Order matters: total response to the same flight (paper's Media 2 settings, K = 0.3)")
    print(order_comparison().to_string(index=False))
    path = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(__file__).with_name("mmm-practice-synthetic.csv")
    if not path.exists():
        sys.exit(f"\nPractice CSV not found at {path}. Download it from /downloads/mmm/mmm-practice-synthetic.csv")
    table, r2, (t0, t1) = fit_practice(path)
    print(f"\nIllustrative least-squares fit on {path.name} (invented data), R^2 = {r2:.3f}")
    print(f"ROAS/mROAS change period: weeks {t0 + 1}-{t1 + 1}, counted through the 12 weeks after it")
    print(table.to_string(index=False))
    edges = [c for c, r in table.set_index("channel").iterrows() if r["alpha"] in (0.0, 0.85) or r["K_scaled"] in (0.2, 10.0)]
    if edges:
        print("\nNote: " + ", ".join(edges) + " landed on the edge of the parameter grid. On this practice CSV that is the"
              "\nexpected answer: it was generated with no carryover and no saturation (true ROAS 1.4 / 0.8 / 0.6), so the"
              "\nfit pushes alpha to zero and K to its upper bound. On real data, a parameter stuck on a bound usually means"
              "\nthe data can't pin it down, as in the paper's shampoo case, and a prior or an experiment decides it.")