The on-sale rush and when to spend: instant sell-outs and the timing of media

import os

for v_ in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS"):
    os.environ[v_] = "1"  # nutpie runs chains as threads: a BLAS pool per chain would fight over the cores

import logging
import sys
import time
import warnings

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
import pytensor.tensor as pt
import scipy.sparse as sp
from scipy import stats
from scipy.interpolate import BSpline
from scipy.optimize import minimize
from scipy.special import logsumexp

from pymc_challenges import data

RANDOM_SEED = 105
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")
logging.getLogger("pymc").setLevel(logging.WARNING)
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning, module="arviz")
sys.modules.pop("jax", None)  # arviz leaves a lazy jax stub that makes matplotlib import all of JAX
pd.set_option("display.width", 160)
T_START = time.time()

BLUE, ORANGE, AQUA, PURPLE, RED = "#2a78d6", "#eb6834", "#1baf7a", "#8c5ac8", "#c8384e"
INK, MUTED, LIGHT = "#222222", "#8a8984", "#d9d8d3"
print(f"PyMC {pm.__version__}, ArviZ {az.__version__}")

ALPHA = 0.4                                 # weekly carryover of adstock, cut off at the event
GAMMA = 0.08                                # growth of the log pool per EUR 1k of pre-sale spend


def beta_true(l, kind="varying"):
    """Log-lift per unit of adstock at lead week l: weak in the awareness months, strong near the event."""
    l = np.asarray(l, float)
    return np.full_like(l, 0.10) if kind == "const" else 0.03 + 0.17 * np.exp(-l / 5.0)


def make_events(seed, n_lo=13, n_hi=39, pool_c=1.3, cap_c=4.0):
    """Window length n_e (weeks on sale), popularity v_e, the pool before pre-sale spend, capacity."""
    rs = np.random.default_rng(seed)
    n = rs.integers(n_lo, n_hi + 1, E)
    v = rs.normal(0.0, 1.0, E)
    in_win = LEADW[None, :] < n[:, None]
    size0 = (Y0 * in_win).sum(1).astype(float)
    hsz = np.array([np.median(size0[HOTEL == h]) for h in (0, 1)])[HOTEL]
    return dict(n=n, v=v, in_win=in_win, hsz=hsz, pool0=pool_c * hsz * np.exp(v), cap=np.ceil(cap_c * hsz))


def pre_spend(ev, seed, policy="random", pre_mean=2.0):
    """Announcement spend per event before the on-sale (EUR 1k in total)."""
    rs = np.random.default_rng(seed)
    if policy == "hot":                                    # the big tours get the announcement money
        return pre_mean * np.exp(0.5 * ev["v"]) * rs.lognormal(-0.02, 0.2, E)
    return pre_mean * rs.lognormal(-0.125, 0.5, E)


def add_signal(ev, P, seed, q_sd=0.3):
    """The pool after pre-sale spend, and its noisy external measurement Q (queue, registrations)."""
    rs = np.random.default_rng(seed)
    ev["P"] = P
    ev["pool"] = ev["pool0"] * np.exp(GAMMA * P)
    ev["logQ"] = np.log(ev["pool"]) + rs.normal(0, q_sd, E)
    return ev


SHAPES = ["common ramp", "front-loaded", "back-loaded", "pulsed", "even"]


def shape_matrix(ev, shape):
    """Relative weekly spend by [event, lead week], zero outside the window, mean 1 over the window."""
    k = ev["n"][:, None] - 1 - LEADW[None, :]              # weeks since the on-sale
    if shape == "common ramp":                             # the habit: on-sale burst, low middle, final push
        r = np.full((E, W), 0.5)
        r[k == 0], r[k == 1] = 3.0, 1.5
        r[:, :4] = 2.0
    elif shape == "front-loaded":
        r = np.exp(-k / 8.0)
    elif shape == "back-loaded":
        r = np.exp(-LEADW[None, :] / 8.0).repeat(E, 0)
    elif shape == "pulsed":
        r = ((k // 2) % 2 == 0).astype(float)              # two weeks on, two off
    else:
        r = np.ones((E, W))
    r = np.where(ev["in_win"], r, 0.0)
    return r / (r.sum(1, keepdims=True) / ev["in_win"].sum(1, keepdims=True))


def spend_plan(ev, policy, seed, exp_share=0.0, exp_budget=False, budget_sd=0.5, rescue_k=0.6, week_sd=0.2):
    """Planned in-window spend S[e, l] (EUR 1k a week). 'rescue' budgets fall with the queue; a share of
    events can be put in a timing experiment (random shape), optionally with a random budget as well."""
    rs = np.random.default_rng(seed)
    rx = np.random.default_rng(seed + 50)                  # the experiment's own draws, so the rest of the world
    ex = rx.random(E) < exp_share                          # (budgets, weekly noise) is the same in every design
    shape_x, B_x = rx.choice(SHAPES[1:], E), rx.lognormal(-0.125, 0.5, E)
    Qf = ev["logQ"] - ev["logQ"].mean()
    B = rs.lognormal(0, budget_sd, E) * (np.exp(-rescue_k * Qf) if policy == "rescue" else 1.0)
    B = B / B.mean()
    shape = np.array(["common ramp"] * E, dtype=object)
    shape[ex] = shape_x[ex]
    if exp_budget:
        B = np.where(ex, B_x, B)
    S = np.zeros((E, W))
    for s in SHAPES:
        if (shape == s).any():
            S[shape == s] = shape_matrix(ev, s)[shape == s]
    S = B[:, None] * S * rs.lognormal(-week_sd ** 2 / 2, week_sd, (E, W)) * ev["in_win"]
    return dict(S=S, B=B, shape=shape)


def world(seed, policy="random", pre="random", exp_share=0.0, exp_budget=False, budget_sd=0.5, week_sd=0.2, **kw):
    ev = make_events(seed, **kw)
    ev = add_signal(ev, pre_spend(ev, seed + 1, pre), seed + 2)
    return ev, spend_plan(ev, policy, seed + 3, exp_share, exp_budget, budget_sd, week_sd=week_sd)


def adstock_t(x):
    """Geometric carryover along lead weeks in time order (lead W-1 first)."""
    out, a = np.zeros_like(x), np.zeros(x.shape[0])
    for l in range(W - 1, -1, -1):
        a = x[:, l] + ALPHA * a
        out[:, l] = a
    return out


def simulate(ev, plan, seed, capped=True, stop=True, reactive=False, beta_kind="varying", phi=0.0, tilt=0.6,
             late_pop=0.3):
    """One world in weekly bins. Late build: the real hotel bookings scaled by popularity (Poisson copies or
    binomial thinning) and tilted toward walk-up for slow sellers; rush: Poisson(pool KER[k]); media
    multiplies both by e^{beta(l) A}. Capacity censors the sell-out week; spend may stop there. reactive may be a
    per-event mask: events in an experiment hold their assigned schedule."""
    ro = np.random.default_rng(seed)
    pace = np.broadcast_to(np.asarray(reactive, bool), (E,))
    U1, U2, U3 = ro.random((E, W)), ro.random((E, W)), ro.random((E, W))
    S, inw, n, v = plan["S"], ev["in_win"], ev["n"], ev["v"][:, None]
    f = np.exp(late_pop * v - tilt * v * (np.exp(-LEADW[None, :] / 4.0) - 0.15))
    base = Y0.astype(float)
    late = np.where(f >= 1, base + stats.poisson.ppf(U1, base * (f - 1)), stats.binom.ppf(U1, Y0, np.minimum(f, 1)))
    late = late * inw
    kk = n[:, None] - 1 - LEADW[None, :]
    rush_rate = np.where((kk >= 0) & (kk < NK), ev["pool"][:, None] * KER[np.clip(kk, 0, NK - 1)], 0.0)
    bt = beta_true(LEADW, beta_kind)
    cap = ev["cap"] if capped else np.full(E, np.inf)
    rem, sold, a, cum, pulled = cap.copy(), np.zeros(E, bool), np.zeros(E), np.zeros(E), np.zeros(E)
    y, A, spend, rush_y = (np.zeros((E, W)) for _ in range(4))
    cens, live = np.zeros((E, W), bool), np.zeros((E, W), bool)
    if pace.any():                                      # pacing target: the planner's expected cumulative sales
        exp_w = (np.exp(ev["logQ"])[:, None] * np.where((kk >= 0) & (kk < NK), KER[np.clip(kk, 0, NK - 1)], 0)
                 + (ev["hsz"] / inw.sum(1))[:, None] * inw)
        target = np.cumsum(exp_w[:, ::-1], 1)[:, ::-1]
    for l in range(W - 1, -1, -1):
        sd = S[:, l].copy()
        if pace.any():                                  # behind target: spend more
            prev = target[:, l + 1] if l + 1 < W else np.zeros(E)
            g_ = np.where(pace & (prev > 0), np.clip((prev - cum) / np.maximum(prev, 1.0), -0.5, 0.5), 0.0)
            sd *= np.exp(2.0 * g_)
        if stop:
            sd = np.where(sold, 0.0, sd)
        a = sd + ALPHA * a
        extra = stats.poisson.ppf(U2[:, l], late[:, l] * np.expm1(bt[l] * a))
        if phi > 0:                                     # a share phi of media buyers would have bought later anyway
            take = ro.binomial(extra.astype(np.int64), phi)
            for e in np.flatnonzero(take):
                fut = late[e, :l].astype(np.int64)
                n_ = min(take[e], fut.sum())
                if n_ > 0:
                    late[e, :l] -= ro.multivariate_hypergeometric(fut, n_)
                    pulled[e] += n_
        r_ = stats.poisson.ppf(U3[:, l], rush_rate[:, l] * np.exp(bt[l] * a))
        yd = (late[:, l] + extra + r_) * inw[:, l]
        live[:, l] = inw[:, l] & ~sold
        hit = live[:, l] & (yd >= rem)
        yd = np.where(live[:, l], np.minimum(yd, rem), 0.0)
        cens[:, l] = hit
        rem -= yd
        cum += yd
        sold |= hit
        y[:, l], A[:, l], spend[:, l] = yd, a * inw[:, l], sd * inw[:, l]
        rush_y[:, l] = np.minimum(r_, yd) * live[:, l]        # rush tickets actually sold (capped at the sell-out)
    so_lead = np.where(cens.any(1), np.argmax(cens, 1), -1)
    return dict(y=y, A=A, spend=spend, cens=cens, live=live, sold=sold, rush_y=rush_y, pulled=pulled,
                so_k=np.where(so_lead >= 0, n - 1 - so_lead, -1), total=y.sum(1),
                rem=cap[:, None] - (np.cumsum(y[:, ::-1], 1)[:, ::-1] - y))


def panel(ev, D, mask=None):
    """Flatten the event x lead-week cells that were on sale (or those in mask)."""
    e, l = np.nonzero(D["live"] if mask is None else mask)
    return dict(y=D["y"][e, l], A=D["A"][e, l], cens=D["cens"][e, l], rem=D["rem"][e, l], ev=e, lead=l,
                k=ev["n"][e] - 1 - l, cal=CALW[e, l])

ev_d, plan_d = world(1, policy="rescue")
D_d = simulate(ev_d, plan_d, 4)
print(f"demo world: {D_d['sold'].mean():.0%} of events sell out, {np.mean(D_d['so_k'] == 0):.0%} in the on-sale week; "
      f"the rush is {D_d['rush_y'].sum() / D_d['y'].sum():.0%} of tickets sold")
print(f"capacity: City {ev_d['cap'][0]:.0f}, Resort {ev_d['cap'][-1]:.0f} tickets; correlation of log Q with the "
      f"log pool {np.corrcoef(ev_d['logQ'], np.log(ev_d['pool']))[0, 1]:.2f}; with log budget "
      f"{np.corrcoef(ev_d['logQ'], np.log(plan_d['B']))[0, 1]:.2f}")
kk_d = ev_d["n"][:, None] - 1 - LEADW[None, :]
fig, axes = plt.subplots(2, 3, figsize=(16, 7.5))
pick = [int(np.flatnonzero((D_d["so_k"] == 0) & (ev_d["n"] > 25))[0]),
        int(np.flatnonzero(D_d["sold"] & (D_d["so_k"] > 8) & (ev_d["n"] > 25))[0]),
        int(np.flatnonzero(~D_d["sold"] & (ev_d["v"] < -0.8) & (ev_d["n"] > 25))[0])]
for c, (e, lab) in enumerate(zip(pick, ["a hot event", "a middle event", "a slow seller"])):
    ax = axes[0, c]
    n_ = ev_d["n"][e]
    ax.bar(LEADW[:n_], D_d["y"][e, :n_] - D_d["rush_y"][e, :n_], color=LIGHT, label="late build")
    ax.bar(LEADW[:n_], D_d["rush_y"][e, :n_], bottom=D_d["y"][e, :n_] - D_d["rush_y"][e, :n_], color=ORANGE,
           label="on-sale rush")
    ax.axvline(n_ - 1, color=RED, lw=1, ls=":")
    if D_d["sold"][e]:
        ax.axvline(n_ - 1 - D_d["so_k"][e], color=RED, lw=1, ls="--")
    ax.set(xlim=(W, -1), ylabel="tickets a week")
    ax.set_title(f"{lab}: {n_} weeks on sale, pool {ev_d['pool'][e]:.0f}\n{D_d['total'][e]:.0f} of {ev_d['cap'][e]:.0f} "
                 "sold" + (f", sold out in on-sale week {D_d['so_k'][e]}" if D_d["sold"][e] else ""), fontsize=10)
    ax = axes[1, c]
    ax.plot(LEADW[:n_], plan_d["S"][e, :n_], color=MUTED, lw=0.9, label="planned spend")
    ax.plot(LEADW[:n_], D_d["spend"][e, :n_], color=ORANGE, lw=1.3, label="actual spend")
    ax.plot(LEADW[:n_], D_d["A"][e, :n_], color=PURPLE, lw=1.3, label="adstock")
    ax.set(xlim=(W, -1), xlabel="lead week (weeks before the event)", ylabel="EUR 1k a week")
axes[0, 0].legend(fontsize=8)
axes[1, 0].legend(fontsize=8);

NKH = 6                                       # rush weeks 0-5 since the on-sale


def dummies(codes, n=None, drop_first=False):
    codes = np.asarray(codes)
    if n is None:
        _, codes = np.unique(codes, return_inverse=True)
        n = codes.max() + 1
    M = sp.csr_matrix((np.ones(len(codes)), (np.arange(len(codes)), codes)), shape=(len(codes), n))
    return M[:, 1:] if drop_first else M


BANDS = [(0, 3), (4, 11), (12, 38)]


def fit_fast(P, model="pool", beta="const", censor=True, event="fe", tau=(1.0, 1.0), mundlak=False, rows=None,
             b_ridge_sd=2.0, qcov=None, return_all=False):
    """Poisson ML on weekly cells (L-BFGS on a sparse design), censored sell-out weeks optional.
    'lead':   log mu = a_e + f(lead) + beta A
    '3clock': log mu = a_e + f(lead) + h(k) + g(cal) + beta A
    'pool':   mu = [exp(a_e + f(lead) + g(cal)) + 1{k < 6} exp(b_e + h(k))] e^{beta A}
    event 'fe' (free effects; a ridge of sd 2 on b_e) or 're' (Gaussian effects of sd tau: a posterior mode)."""
    r = np.ones(len(P["y"]), bool) if rows is None else rows
    y, A, ev, lead, k, cal = (P[c][r] for c in ("y", "A", "ev", "lead", "k", "cal"))
    c = P["cens"][r] if censor else np.zeros(len(y), bool)
    _, evi = np.unique(ev, return_inverse=True)
    nE = evi.max() + 1
    Bb = np.column_stack([((lead >= a) & (lead <= b)).astype(float) for a, b in BANDS]) if beta == "bands" \
        else np.ones((len(y), 1))
    Xb, kb = Bb * A[:, None], Bb.shape[1]
    blocks1 = [dummies(evi, nE), dummies(lead, W, drop_first=True)]
    if model == "3clock":
        blocks1.append(dummies(np.minimum(k, NKH), NKH + 1)[:, :NKH])
    if model in ("3clock", "pool"):
        blocks1.append(dummies(cal, cal.max() + 1, drop_first=True))
    if event == "re" or mundlak:
        M_ev = np.log(np.maximum(np.bincount(evi, weights=A, minlength=nE) / np.bincount(evi, minlength=nE), 1e-3))
        M_ev -= M_ev.mean()
    if event == "re":
        blocks1.append(sp.csr_matrix(np.ones((len(y), 1))))
        if mundlak:
            blocks1.append(sp.csr_matrix(M_ev[evi][:, None]))
    X1 = sp.hstack(blocks1).tocsr()
    n1 = X1.shape[1]
    pen1 = np.full(n1, 1e-3)
    pen1[:nE] = 1e-4 if event == "fe" else 0.5 / tau[0] ** 2
    if event == "re":
        pen1[n1 - 1 - int(mundlak):] = 0.0
    rr = k < NKH
    if model == "pool":
        blocks2 = [dummies(evi, nE), dummies(np.minimum(k, NKH - 1), NKH, drop_first=True),
                   sp.csr_matrix(np.ones((len(y), 1)))]
        if mundlak:
            blocks2.append(sp.csr_matrix(M_ev[evi][:, None]))
        if qcov is not None:
            blocks2.append(sp.csr_matrix(qcov[r][:, None]))
        X2 = (sp.diags(rr.astype(float)) @ sp.hstack(blocks2)).tocsr()
        n2 = X2.shape[1]
        pen2 = np.full(n2, 1e-3)
        pen2[:nE] = 0.5 / (b_ridge_sd if event == "fe" else tau[1]) ** 2
        pen2[nE + NKH - 1:] = 0.0
    else:
        X2, n2, pen2 = None, 0, np.zeros(0)
    X1T, X2T = X1.T.tocsr(), (X2.T.tocsr() if X2 is not None else None)
    yc, pen = y[c], np.concatenate([np.zeros(kb), pen1, pen2])

    def nll(th):
        m = np.exp(np.clip(Xb @ th[:kb], -20, 20))
        e1 = np.exp(np.clip(X1 @ th[kb:kb + n1], -30, 20))
        e2 = np.where(rr, np.exp(np.clip(X2 @ th[kb + n1:], -30, 20)), 0.0) if X2 is not None else 0.0
        mu = (e1 + e2) * m
        ll, gmu = y * np.log(mu) - mu, y / mu - 1.0
        if c.any():                                        # log P(Y >= R) and its derivative
            lsf = stats.poisson.logsf(yc - 1, mu[c])
            with np.errstate(all="ignore"):
                hz = np.exp(stats.poisson.logpmf(yc - 1, mu[c]) - lsf)
            ll[c], gmu[c] = lsf, np.where(np.isfinite(hz), hz, yc / mu[c] - 1.0)
        g = [Xb.T @ (gmu * mu), X1T @ (gmu * e1 * m)] + ([X2T @ (gmu * e2 * m)] if X2 is not None else [])
        return -ll.sum() + (pen * th * th).sum(), -np.concatenate(g) + 2 * pen * th

    th0 = np.zeros(kb + n1 + n2)
    lm = np.log(np.bincount(evi, weights=y, minlength=nE) / np.bincount(evi, minlength=nE) + 0.1)
    if model == "pool":
        th0[kb:kb + nE] = lm - 1.0 - (lm.mean() if event == "re" else 0)
        th0[kb + n1:kb + n1 + nE] = lm + 0.5 - (lm + 0.5).mean()
        th0[kb + n1 + nE + NKH - 1] = (lm + 0.5).mean() - (np.mean(qcov[r]) if qcov is not None else 0.0)
        if qcov is not None:
            th0[-1] = 1.0
        th0[kb + n1 + nE:kb + n1 + nE + NKH - 1] = np.log(KER[1:] / KER[0])   # start at the real rush shape
    else:
        th0[kb:kb + nE] = lm - (lm.mean() if event == "re" else 0)
    if event == "re":
        th0[kb + n1 - 1 - int(mundlak)] = lm.mean() - (1.0 if model == "pool" else 0)
    res = minimize(nll, th0, jac=True, method="L-BFGS-B", options=dict(maxiter=20000, gtol=1e-6))
    if return_all:
        return dict(x=res.x, kb=kb, n1=n1, nE=nE, ev_u=np.unique(ev))
    return res.x[:kb] if kb > 1 else float(res.x[0])


def tau_from_fe(P, model):
    """Spread of the fixed effects: the sd for the random-effects fits."""
    o = fit_fast(P, model=model, return_all=True)
    kb, n1, nE = o["kb"], o["n1"], o["nE"]
    out = [np.std(o["x"][kb:kb + nE]), np.std(o["x"][kb + n1:kb + n1 + nE]) if model == "pool" else 1.0]
    return tuple(max(float(v), 0.1) for v in out)


def rush_hat(o, P, qcov=None):
    """Expected rush buyers per event at zero spend, from a 'pool' fit."""
    t2 = o["x"][o["kb"] + o["n1"]:]
    nE = o["nE"]
    eta = t2[:nE] + t2[nE + NKH - 1]
    if qcov is not None:
        _, first = np.unique(P["ev"], return_index=True)
        eta = eta + t2[-1] * qcov[first]
    return np.exp(eta) * (1 + np.exp(t2[nE:nE + NKH - 1]).sum())


R_REP = 2
SC = []


def rec(part, scen, est, val, truth):
    SC.append(dict(part=part, scenario=scen, estimator=est, value=val, truth=truth))


t0 = time.time()
for r in range(R_REP):
    for scen, wkw, flat in [("windows 3-9 months, common ramp", {}, False),
                            ("every window 26 weeks, common ramp", dict(n_lo=26, n_hi=26), False),
                            ("windows 3-9 months, even spend", {}, True)]:
        ev, plan = world(100 + r, **wkw)
        if flat:
            plan = dict(plan, S=plan["B"][:, None] * shape_matrix(ev, "even")
                        * np.random.default_rng(5).lognormal(-0.02, 0.2, (E, W)) * ev["in_win"])
        D = simulate(ev, plan, 200 + r, capped=False, stop=False, beta_kind="const")
        P = panel(ev, D)
        for m, lab in [("lead", "lead time only"), ("3clock", "three clocks, log-additive"), ("pool", "pool model")]:
            rec("C", scen, lab, fit_fast(P, model=m, censor=False), 0.10)
tab_c = pd.DataFrame(SC).groupby(["scenario", "estimator"], sort=False).value.agg(["mean", "min", "max"])
print(f"true beta 0.10 (no capacity); {R_REP} worlds per row, {time.time() - t0:.0f} s")
print(tab_c.round(3).to_string())

t0 = time.time()
dem = []
for r in range(R_REP):
    ev, plan = world(300 + r)
    D = simulate(ev, plan, 400 + r, capped=True, stop=True, beta_kind="const")
    Pall = panel(ev, D, mask=ev["in_win"])                       # every week of the window, sold out or not
    P = panel(ev, D)
    inst = D["so_k"] == 0
    scn = "capacity, spend stops at sell-out"
    rec("D", scn, "every week kept, no censoring", fit_fast(Pall, censor=False), 0.10)
    rec("D", scn, "weeks after the sell-out dropped, sell-out week censored", fit_fast(P), 0.10)
    rec("D", scn, "... and the instant sell-outs dropped", fit_fast(P, rows=~inst[P["ev"]]), 0.10)
    q = ev["logQ"][P["ev"]]
    for lab, kw in [("no queue signal", dict(b_ridge_sd=1.0)), ("queue as a covariate", dict(b_ridge_sd=0.35, qcov=q))]:
        o = fit_fast(P, return_all=True, **kw)
        ratio = rush_hat(o, P, qcov=kw.get("qcov")) / ev["pool"][o["ev_u"]]
        ii = inst[o["ev_u"]]
        dem.append(dict(world=r, model=lab, events="instant sell-outs", median_ratio=np.median(ratio[ii]),
                        median_abs_log_error=np.median(np.abs(np.log(ratio[ii])))))
        dem.append(dict(world=r, model=lab, events="the rest", median_ratio=np.median(ratio[~ii]),
                        median_abs_log_error=np.median(np.abs(np.log(ratio[~ii])))))
    sold6 = (D["y"] * ((ev["n"][:, None] - 1 - LEADW[None, :]) < 6)).sum(1)
    dem.append(dict(world=r, model="tickets sold = demand", events="instant sell-outs",
                    median_ratio=np.median(sold6[inst] / ev["pool"][inst]),
                    median_abs_log_error=np.median(np.abs(np.log(sold6[inst] / ev["pool"][inst])))))
    print(f"world {r}: {D['sold'].mean():.0%} sell out, {inst.mean():.0%} in the on-sale week")
tab_d = pd.DataFrame([s for s in SC if s["part"] == "D"]).groupby("estimator", sort=False).value.agg(["mean", "min", "max"])
print(f"true beta 0.10; {time.time() - t0:.0f} s")
print(tab_d.round(3).to_string())
print("estimated over true pool (demand) per event, median over events, mean of the worlds:")
print(pd.DataFrame(dem).groupby(["events", "model"], sort=False)[["median_ratio", "median_abs_log_error"]].mean()
      .round(2).to_string())

KN = np.array([0, 0, 0, 0, 3, 9, 20, 38, 38, 38, 38], float)
BS = BSpline.design_matrix(np.arange(W, dtype=float), KN, 3).toarray()   # 7 cubic B-splines on lead weeks 0-38
BT = beta_true(LEADW)


def pool_model(ev, D, EV):
    """Censored negative-binomial pool model on the weekly cells of events EV."""
    e_, l_ = np.nonzero(D["live"][EV])
    ge = EV[e_]
    y, A, cen, rem = D["y"][ge, l_], D["A"][ge, l_], D["cens"][ge, l_], D["rem"][ge, l_]
    k_ = ev["n"][ge] - 1 - l_
    cal_u, cal_i = np.unique(CALW[ge, l_], return_inverse=True)
    ncal, nE = len(cal_u), len(EV)
    qc = ev["logQ"][EV] - ev["logQ"][EV].mean()
    Pc = ev["P"][EV] - ev["P"].mean()
    hot = HOTEL[EV]
    tt = np.arange(ncal) - (ncal - 1) / 2
    Qm = np.column_stack([np.ones(ncal), tt])
    PD = np.eye(ncal) - Qm @ np.linalg.pinv(Qm)                 # detrended calendar walk, as in E104
    unc, cc = np.flatnonzero(~cen), np.flatnonzero(cen)
    rr, kk = (k_ < NKH).astype(float), np.minimum(k_, NKH - 1)
    coords = dict(event=np.arange(nE), leadweek=LEADW, k=np.arange(1, NKH), calweek=np.arange(ncal),
                  step=np.arange(1, BS.shape[1]), hotel=HOTELS)
    with pm.Model(coords=coords) as m:
        mu_a = pm.Normal("mu_a", np.log(np.mean(y[k_ >= NKH]) + 0.1), 1.0, dims="hotel")
        c_aq = pm.Normal("c_aq", 0.0, 1.0)
        sd_a = pm.HalfNormal("sd_a", 0.5)
        z_a = pm.Normal("z_a", 0.0, 1.0, dims="event")
        a_ev = mu_a[hot] + c_aq * qc + sd_a * z_a
        curve0 = pm.ZeroSumNormal("curve0", 1.0, dims="leadweek")
        sd_hc = pm.HalfNormal("sd_hc", 0.3)                   # how far each hotel's curve strays from the common one
        z_hc = pm.ZeroSumNormal("z_hc", 1.0, dims=("hotel", "leadweek"), n_zerosum_axes=2)
        curve = pm.Deterministic("curve", curve0[None, :] + sd_hc * z_hc, dims=("hotel", "leadweek"))
        sd_cal = pm.HalfNormal("sd_cal", 0.2)
        cal_innov = pm.Normal("cal_innov", 0.0, 1.0, dims="calweek")
        calv = pt.dot(PD, pt.cumsum(cal_innov) * sd_cal)
        mu_b = pm.Normal("mu_b", np.log(np.mean(y[k_ == 0]) + 1), 1.5, dims="hotel")
        gamma = pm.Normal("gamma", 0.0, 0.5)
        sd_b = pm.HalfNormal("sd_b", 1.0)
        z_b = pm.Normal("z_b", 0.0, 1.0, dims="event")
        m_b = mu_b[hot] + gamma * Pc
        q0 = pm.Normal("q0", -np.log(np.mean(y[k_ == 0]) + 1), 1.5)
        s_q = pm.HalfNormal("s_q", 0.5)
        v_ = sd_b ** 2 + s_q ** 2                               # the queue's marginal, then b_e given the queue
        pm.Normal("logQ", m_b + q0, pt.sqrt(v_), observed=qc)
        b_ev = m_b + sd_b ** 2 / v_ * (qc - q0 - m_b) + sd_b * s_q / pt.sqrt(v_) * z_b
        h = pm.Normal("h", -1.5, 1.5, dims="k")
        hk = pt.concatenate([pt.zeros(1), h])
        c0 = pm.Normal("c0", 0.1, 0.15)
        sd_c = pm.HalfNormal("sd_c", 0.05)
        zc = pm.Normal("zc", 0.0, 1.0, dims="step")
        coef = c0 + pt.concatenate([pt.zeros(1), pt.cumsum(sd_c * zc)])   # a random walk over the spline weights
        beta_l = pm.Deterministic("beta_l", pt.dot(BS, coef), dims="leadweek")
        alpha_r = pm.Gamma("alpha_r", 2.0, 0.05)              # dispersion of the on-sale weeks
        alpha_l = pm.Gamma("alpha_l", 2.0, 0.2)               # ... and of the late build
        al_c = pt.switch(k_ < NKH, alpha_r, alpha_l)
        mu = (pt.exp(a_ev[e_] + curve[hot[e_], l_] + calv[cal_i]) + rr * pt.exp(b_ev[e_] + hk[kk])) * pt.exp(beta_l[l_] * A)
        pm.NegativeBinomial("y", mu=mu[unc], alpha=al_c[unc], observed=y[unc])
        pm.Censored("y_soldout", pm.NegativeBinomial.dist(mu=mu[cc], alpha=al_c[cc]), lower=None,
                    upper=rem[cc], observed=rem[cc])
    info = dict(y=y, k=k_, l=l_, e=e_, unc=unc, cc=cc, n_inst=int(np.sum(cen & (k_ == 0))), qmean=ev["logQ"][EV].mean(),
                Pmean=ev["P"].mean())
    return m, info


VARS = ["mu_a", "c_aq", "sd_a", "z_a", "curve0", "sd_hc", "z_hc", "sd_cal", "cal_innov", "mu_b", "gamma", "sd_b", "z_b", "q0", "s_q",
        "h", "c0", "sd_c", "zc", "alpha_r", "alpha_l", "beta_l"]


INIT = {"alpha_l": 1.2, "alpha_r": 50.0, "s_q": 0.3, "sd_b": 1.0, "sd_hc": 0.2}   # a sensible start for the scales


def fit_report(name, m, info, ev, D, EV, draws=1000):
    t0 = time.time()
    with m:
        idata = pm.sample(draws=draws, target_accept=0.9, random_seed=RANDOM_SEED, progressbar=False,
                          initvals=INIT)
    t = time.time() - t0
    st, po = idata.sample_stats, idata.posterior
    rh = {v: float(az.rhat(po[v]).max()) for v in VARS}
    es = min(min(float(az.ess(po[v], method=k).min()) for k in ("bulk", "tail")) for v in VARS)
    bl = po["beta_l"].to_numpy().reshape(-1, W)
    lo, hi = np.quantile(bl, [0.05, 0.95], axis=0)
    worst = sorted(rh, key=rh.get)[-3:][::-1]
    print(f"{name}: {len(EV)} events, {len(info['y']):,} event-weeks, {len(info['cc'])} censored sell-out weeks "
          f"({info['n_inst']} in the on-sale week); {t:.0f} s, {po.attrs.get('tuning_steps')} warm-up steps, "
          f"{int(st['diverging'].sum())} divergences, "
          f"max r_hat {max(rh.values()):.3f} ({', '.join(worst)}), min ESS {es:.0f}, max depth {int(st['depth'].max())}")
    print("   beta(lead week), mean (90% interval) against the truth: " + "; ".join(
        f"{l}: {bl[:, l].mean():.3f} ({lo[l]:.3f}-{hi[l]:.3f}) vs {BT[l]:.3f}" for l in (0, 2, 5, 10, 20)))
    print(f"   90% band covers the truth at {np.mean((lo <= BT) & (BT <= hi)):.0%} of lead weeks")
    for v in ("gamma", "c_aq", "sd_a", "sd_b", "s_q", "alpha_r", "alpha_l"):
        x = po[v].to_numpy().ravel()
        print(f"   {v}: {x.mean():.3f} ({np.quantile(x, 0.05):.3f}-{np.quantile(x, 0.95):.3f})", end="")
    print()
    sub = az.extract(idata, var_names=["mu_a", "c_aq", "sd_a", "mu_b", "gamma", "sd_b", "q0", "s_q", "h", "curve",
                                       "alpha_r", "alpha_l", "beta_l"], num_samples=400, random_seed=1)
    post = {v: sub[v].values for v in sub.data_vars}
    post.update(beta_draws=post.pop("beta_l"), qmean=info["qmean"], Pmean=info["Pmean"])
    return idata, post, bl


# observational world: common ramp, rescue budgets, reactive pacing, the walk-up tilt
ev_o, plan_o = world(700, policy="rescue", budget_sd=0.2, week_sd=0.05)
D_o = simulate(ev_o, plan_o, 800, reactive=True)
rs_ = np.random.default_rng(7)
EV_O = np.sort(np.concatenate([rs_.choice(np.flatnonzero(HOTEL == h), 80, replace=False) for h in (0, 1)]))
m_obs, info_o = pool_model(ev_o, D_o, EV_O)
with m_obs:
    prior = pm.sample_prior_predictive(300, random_seed=RANDOM_SEED)
pb = prior.prior["beta_l"].to_numpy().reshape(-1, W)
pp_y = prior.prior_predictive["y"].to_numpy().reshape(300, -1)
print("prior beta(lead 0) 5/50/95%: " + " / ".join(f"{v:.2f}" for v in np.percentile(pb[:, 0], [5, 50, 95]))
      + "; beta(lead 20): " + " / ".join(f"{v:.2f}" for v in np.percentile(pb[:, 20], [5, 50, 95])))
print("prior predictive weekly tickets (uncensored weeks) 5/50/95%: " + " / ".join(
    f"{v:,.0f}" for v in np.percentile(pp_y, [5, 50, 95])) + "; observed " + " / ".join(
    f"{v:.0f}" for v in np.percentile(info_o["y"][info_o["unc"]], [5, 50, 95])))