import logging
import time
import warnings
from dataclasses import dataclass

import arviz as az
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
import pytensor
import pytensor.tensor as pt
from IPython.display import display

from pymc_challenges import data

RANDOM_SEED = 62
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")
logging.getLogger("pymc").setLevel(logging.WARNING)  # many small fits: no sampler banner each
warnings.filterwarnings("ignore", category=RuntimeWarning, module="arviz")  # NaN r_hat of constants
BLUE, ORANGE, AQUA, GREY, PURPLE, RED = "#2a78d6", "#eb6834", "#1baf7a", "#8a8a86", "#8c5ac8", "#c8384e"
print(f"PyMC {pm.__version__}, ArviZ {az.__version__}, PyTensor {pytensor.__version__}")

@dataclass
class Question:
    asked: str        # the question in the stakeholder's words
    quantity: str     # the number that answers it
    population: str   # for whom / over what period
    scale: str        # units: absolute, relative, per 100, per $1000...
    decision: str     # what will be done with the answer

    def show(self):
        for k, v in self.__dict__.items():
            print(f"{k:>10}: {v}")


def q(x, probs=(0.05, 0.5, 0.95)):
    """Posterior quantiles of a 1-D array of draws, rounded for printing."""
    return np.round(np.quantile(x, probs), 2)

QA = Question(
    asked="Which channel brings more new customers per dollar, Google or Meta? "
          "Where should next quarter's extra budget go?",
    quantity="ROI of each channel: incremental new customers per $1000, and their difference",
    population="this brand, Aug 2021 - May 2024, at the spend levels it actually ran",
    scale="new customers per $1000 (CAC = 1000 / ROI dollars)",
    decision="shift budget between Google and Meta",
)
QA.show()

roi_a = az.extract(idata_a, var_names=["roi"]).values  # (channel, sample)
print("ROI Google  (customers per $1000):", q(roi_a[0]))
print("ROI Meta                         :", q(roi_a[1]))
print(f"P(Google ROI > Meta ROI) = {(roi_a[0] > roi_a[1]).mean():.3f}")

roi_p = prior_a.prior["roi"].values.reshape(-1, 2)
tot_a = az.extract(idata_a, var_names=["roi_total"]).values
fig, axes = plt.subplots(1, 2, figsize=(11, 4))
axes[0].scatter(roi_p[:1500, 0], roi_p[:1500, 1], s=3, color=GREY, alpha=0.3, label="prior")
axes[0].scatter(roi_a[0, ::4], roi_a[1, ::4], s=3, color=BLUE, alpha=0.4, label="posterior")
axes[0].plot([0, 80], [0, 80], color="k", lw=0.8, ls="--")
axes[0].set(xlim=(0, 80), ylim=(0, 80), xlabel="ROI Google", ylabel="ROI Meta",
            title=f"Posterior correlation {np.corrcoef(roi_a)[0, 1]:.2f}")
axes[0].legend(markerscale=4)
for x_, col, lab in [(roi_a[0], BLUE, "Google"), (roi_a[1], ORANGE, "Meta"), (tot_a, "k", "total")]:
    axes[1].hist(x_, bins=60, density=True, histtype="step", color=col, lw=1.5, label=lab)
axes[1].set(xlabel="customers per $1000", title="Per-channel ROI and total ROI")
axes[1].legend();

def contraction(prior_draws, post_draws):
    return 1 - np.var(post_draws) / np.var(prior_draws)


pri = {k: prior_a.prior[k].values.reshape(-1, *prior_a.prior[k].shape[2:]) for k in
       ["roi", "roi_diff", "roi_total", "mroi"]}
post = {k: az.extract(idata_a, var_names=[k]).values for k in ["roi_diff", "roi_total", "mroi"]}
print(pd.Series({
    "ROI Google": contraction(pri["roi"][:, 0], roi_a[0]),
    "ROI Meta": contraction(pri["roi"][:, 1], roi_a[1]),
    "ROI Google - Meta": contraction(pri["roi_diff"], post["roi_diff"]),
    "ROI total": contraction(pri["roi_total"], post["roi_total"]),
    "marginal ROI (+10%)": contraction(pri["mroi"], post["mroi"]),
}, name="contraction").round(3).to_string())

PRIORS = {"neutral (600, 600)": (600.0, 600.0), "Google-leaning (1200, 300)": (1200.0, 300.0),
          "Meta-leaning (300, 1200)": (300.0, 1200.0)}
fits_a = {"neutral (600, 600)": idata_a}
for name, sd in list(PRIORS.items())[1:]:
    with build_mmm(beta_sd=sd):
        fits_a[name] = pm.sample(random_seed=RANDOM_SEED, progressbar=False, var_names=KEEP)

rows = []
for name, idt in fits_a.items():
    r = az.extract(idt, var_names=["roi"]).values
    rows.append({"prior on beta": name,
                 "ROI Google": f"{np.median(r[0]):.1f}", "ROI Meta": f"{np.median(r[1]):.1f}",
                 "P(Google > Meta)": round(float((r[0] > r[1]).mean()), 3),
                 "ROI total (5-95%)": str(q(az.extract(idt, var_names=["roi_total"]).values, (0.05, 0.95))),
                 "divergences": int(idt.sample_stats["diverging"].sum())})
sens_a = pd.DataFrame(rows).set_index("prior on beta")
sens_a

fig, ax = plt.subplots(figsize=(6.5, 5))
for (name, idt), col in zip(fits_a.items(), [BLUE, AQUA, ORANGE]):
    r = az.extract(idt, var_names=["roi"]).values
    ax.scatter(r[0, ::4], r[1, ::4], s=3, alpha=0.35, color=col, label=name)
ax.plot([0, 70], [0, 70], color="k", lw=0.8, ls="--")
ax.text(38, 45, "Google = Meta", fontsize=9)
ax.set(xlim=(10, 70), ylim=(10, 50), xlabel="ROI Google (customers per $1000)", ylabel="ROI Meta",
       title="Each prior picks a different point on the same ridge")
ax.legend(markerscale=5, fontsize=9);

with build_mmm(baseline="rw") as mmm_rw:
    idata_rw = pm.sample(random_seed=RANDOM_SEED, progressbar=False, target_accept=0.9)
print("divergences:", int(idata_rw.sample_stats["diverging"].sum()))
az.summary(idata_rw, var_names=["roi", "roi_total", "mroi", "s_rw", "alpha"], round_to=3)

base_share = {k: float(az.extract(i, var_names=["base"]).values.mean(axis=1).sum() / y.sum())
              for k, i in [("linear", idata_a), ("random walk", idata_rw)]}
print("share of new customers the baseline explains:", {k: round(v, 2) for k, v in base_share.items()})
print("ROI total, linear baseline      :", q(tot_a))
print("ROI total, random-walk baseline :", q(az.extract(idata_rw, var_names=["roi_total"]).values))

loo_a = {}
for name, m, idt in [("linear baseline", mmm, idata_a), ("random-walk baseline", mmm_rw, idata_rw)]:
    ll = idt.copy()
    with m:
        pm.compute_log_likelihood(ll, progressbar=False)
    loo_a[name] = az.loo(ll)
az.compare(loo_a, round_to=1)

truth = dict(theta=np.array([0.30, 0.15]), lam=np.array([0.35, 0.35]),
             beta=np.array([600.0, 2000.0]), b0=3.98, b1=0.17, alpha=30.0)
w_true = (1 - truth["theta"])[:, None] * truth["theta"][:, None] ** np.arange(L_ADSTOCK)
a_true = (lags * (w_true / w_true.sum(axis=1, keepdims=True))[None]).sum(axis=-1)
contrib_true = truth["beta"] * (1 - np.exp(-truth["lam"] * a_true))
roi_true = 1000 * contrib_true.sum(axis=0) / X.sum(axis=0)
mu_true = np.exp(truth["b0"] + truth["b1"] * tt) + contrib_true.sum(axis=1)
rng_fake = np.random.default_rng(3)
y_fake = rng_fake.negative_binomial(truth["alpha"], truth["alpha"] / (truth["alpha"] + mu_true))
lift_obs = roi_true[0] * (1 + 0.10 * rng_fake.normal())
print("true ROI (Google, Meta):", roi_true.round(1), f"| experiment reports Google ROI {lift_obs:.1f} "
      f"+- {0.10 * roi_true[0]:.1f}")

rows, fake_draws = [], {}
for with_lift in (False, True):
    for name, sd in PRIORS.items():
        lift = (0, lift_obs, 0.10 * roi_true[0]) if with_lift else None
        with build_mmm(beta_sd=sd, y_obs=y_fake, lift=lift):
            idf = pm.sample(random_seed=RANDOM_SEED, progressbar=False, var_names=["roi"])
        r = az.extract(idf, var_names=["roi"]).values
        fake_draws[(with_lift, name)] = r
        rows.append({"experiment": "yes" if with_lift else "no", "prior on beta": name,
                     "ROI Google": f"{np.median(r[0]):.1f}", "ROI Meta": f"{np.median(r[1]):.1f}",
                     "P(Google > Meta)": round(float((r[0] > r[1]).mean()), 3),
                     "divergences": int(idf.sample_stats["diverging"].sum())})
fake_tab = pd.DataFrame(rows).set_index(["experiment", "prior on beta"])
fake_tab

fig, axes = plt.subplots(1, 2, figsize=(12, 4.3), sharex=True, sharey=True)
for ax, with_lift in zip(axes, (False, True)):
    for (name, _), col in zip(PRIORS.items(), [BLUE, AQUA, ORANGE]):
        r = fake_draws[(with_lift, name)]
        ax.scatter(r[0, ::4], r[1, ::4], s=3, alpha=0.35, color=col, label=name)
    ax.scatter(*roi_true, marker="*", s=250, color="k", zorder=5)
    ax.annotate("truth", roi_true, xytext=(8, 8), textcoords="offset points", fontsize=10)
    ax.plot([0, 70], [0, 70], color="k", lw=0.8, ls="--")
    ax.set(xlim=(5, 55), ylim=(15, 50), xlabel="ROI Google",
           title="spend history + experiment" if with_lift else "spend history only (fake data)")
axes[0].set_ylabel("ROI Meta")
axes[0].legend(markerscale=5, fontsize=8);

values = np.linspace(20, 120, 201)  # what one new customer is worth to the brand, $
curves = {}
for name, idt in list(fits_a.items()) + [("random-walk baseline", idata_rw)]:
    mcac = 1000 / az.extract(idt, var_names=["mroi"]).values
    curves[name] = [(mcac < v).mean() for v in values]
    print(f"{name:>28}: marginal cost per new customer ($) {q(mcac)}")

fig, ax = plt.subplots(figsize=(8, 4))
for (name, cur), col in zip(curves.items(), [BLUE, AQUA, ORANGE, RED]):
    ax.plot(values, cur, color=col, lw=2, label=name)
v_yes = values[np.argmax(np.min(list(curves.values()), axis=0) > 0.9)]  # all models: P > 0.9
v_no = values[np.argmax(np.max(list(curves.values()), axis=0) > 0.1)]  # below: all P < 0.1
ax.axvspan(v_no, v_yes, color=GREY, alpha=0.2)
ax.text(v_no + 0.6 * (v_yes - v_no), 0.55, "the answer\ndepends on\nthe model", ha="center", fontsize=9)
ax.set(xlabel="value of one new customer to the brand ($)",
       ylabel="P(extra spend pays for itself)",
       title="A decision the unidentified parts do not change - outside the grey band")
ax.legend(fontsize=8, loc="lower right");
print(f"all four models agree (P < 0.1 or P > 0.9) outside ${v_no:.0f} - ${v_yes:.0f}")

audit = []
audit.append(dict(part="A", asked="Is Google or Meta more efficient?",
                  answered="not answerable from this spend history: it follows the prior "
                           f"(P(Google better) {sens_a['P(Google > Meta)'].min():.2f}-"
                           f"{sens_a['P(Google > Meta)'].max():.2f} across three priors)",
                  why="spend on the two channels moved together; the data fix a weighted sum",
                  next_step="geo holdout on one channel (fake-data check: a 10%-SE test on Google "
                            "makes all priors agree)"))
audit.append(dict(part="A", asked="(fallback) total return of paid media",
                  answered="depends on the baseline assumption (median "
                           f"{np.median(az.extract(idata_rw, var_names=['roi_total']).values):.0f} "
                           f"vs {np.median(tot_a):.0f} customers per $1000)",
                  why="slow growth and spend moved together; LOO ranks fit, not attribution",
                  next_step="report both; the same experiment calibrates the total"))
audit.append(dict(part="A", asked="Should the total budget go up?",
                  answered=f"yes if a customer is worth > ${v_yes:.0f}, no if < ${v_no:.0f}, "
                           "undetermined between",
                  why="the decision threshold falls outside the band for most values",
                  next_step="the brand supplies the value of a customer"))

QB = Question(
    asked="How much did underlying bike-share demand grow from 2011 to 2012, beyond weather "
          "and the day of the week? We are sizing next year's fleet.",
    quantity="ratio of the smooth demand level, 2012 average / 2011 average",
    population="Capital Bikeshare, Washington DC, all stations",
    scale="multiplicative growth (1.5 = 50% more rides)",
    decision="fleet size for 2013",
)
QB.show()
data.describe("bike_hour")
bikes = data.load("bike_hour")
bikes["dteday"] = pd.to_datetime(bikes["dteday"])
print(f"{len(bikes):,} hourly rows, {bikes.dteday.min().date()} -> {bikes.dteday.max().date()}")

def growth(level_draws, years):
    return level_draws[years == 2012].mean(axis=0) / level_draws[years == 2011].mean(axis=0)


lvl_day = az.extract(idata_day, var_names=["level"]).values
growth_tab = pd.DataFrame({
    "weekly totals, exact GP": q(growth(np.exp(fw_exact), year_wk)),
    "weekly totals, HSGP": q(growth(np.exp(fw_hsgp), year_wk)),
    "daily totals, HSGP (+ weather, working day)": q(growth(lvl_day, year_day)),
}, index=["5%", "50%", "95%"]).T
growth_tab

fig, ax = plt.subplots(figsize=(11, 4))
ax.plot(day.dteday, day.cnt, color=GREY, lw=0.6, alpha=0.7, label="daily rentals")
lv = np.quantile(lvl_day, [0.05, 0.5, 0.95], axis=1)
ax.fill_between(day.dteday, lv[0], lv[2], color=BLUE, alpha=0.3)
ax.plot(day.dteday, lv[1], color=BLUE, lw=2, label="daily HSGP level (clear, non-working day)")
fe = np.quantile(np.exp(fw_exact) / 7, [0.05, 0.5, 0.95], axis=1)
ax.plot(weekly.index - pd.Timedelta(days=3), fe[1], color=ORANGE, lw=2, ls="--",
        label="weekly exact GP / 7")
ax.set(ylabel="rentals per day", title="Two pivots, one answer to the question that was asked")
ax.legend(fontsize=9);

audit.append(dict(part="B", asked="Growth in underlying demand, 2011 -> 2012",
                  answered=f"about {growth_tab.iloc[2, 1]:.2f}x (90%: {growth_tab.iloc[2, 0]:.2f}-"
                           f"{growth_tab.iloc[2, 2]:.2f}); weekly exact GP agrees",
                  why="exact GP on 17k hours: weeks of computing, GBs per matrix",
                  next_step="gave up hour-level structure (not asked); HSGP checked against exact"))

QC1 = Question(
    asked="Our patients are mostly symptomatic and have had antiretrovirals before. Out of 100 "
          "such patients, how many fewer will have an event within 2 years on ZDV + ddI than on ZDV?",
    quantity="2-year risk difference, ZDV minus ZDV+ddI",
    population="symptomatic, ART-experienced patients (the clinic's case mix)",
    scale="events per 100 patients over 730 days",
    decision="whether switching these patients justifies the extra toxicity of ddI",
)
QC2 = Question(
    asked="Which of the three newer regimens is best?",
    quantity="P(each regimen has the lowest hazard)",
    population="trial population", scale="probability", decision="formulary choice")
QC3 = Question(
    asked="Which kind of patient benefits most from adding ddI?",
    quantity="subgroup-specific effects of ddI, and their ranking",
    population="8 subgroups: symptoms x prior ART x CD4 above/below 340",
    scale="relative (hazard ratio) - or absolute?", decision="whom to prioritise")
QC1.show()

b_draws = az.extract(idata_c, var_names=["b_arm"]).values  # (3, S)
b_draws = np.vstack([np.zeros(b_draws.shape[1]), b_draws])  # ZDV = reference
lb_draws = az.extract(idata_c, var_names=["log_base"]).values
bx_draws = az.extract(idata_c, var_names=["b_x"]).values
days_in_window = np.clip(730 - cuts[:-1], 0, np.diff(cuts))  # exposure to day 730 per window
H0_730 = (np.exp(lb_draws) * days_in_window[:, None]).sum(axis=0)  # baseline cumulative hazard
lin_pred = Xc @ bx_draws  # (patients, S)


def risk_730(j, rows):
    """2-year event risk under arm j for the selected patients: (patients, S)."""
    return 1 - np.exp(-H0_730[None] * np.exp(lin_pred[rows] + b_draws[j][None]))


clinic = ((actg.symptom == 1) & (actg.str2 == 1)).to_numpy()
low_risk = ((actg.symptom == 0) & (actg.str2 == 0) & (actg.cd40 > 350)).to_numpy()
everyone = np.ones(len(actg), bool)
hr = np.exp(b_draws[1])
answers = {
    "hazard ratio ZDV+ddI vs ZDV": q(hr),
    "  ...misread as 'fewer events per 100'": q(100 * (1 - hr)),
    "risk difference per 100, trial population": q(100 * (risk_730(0, everyone) - risk_730(1, everyone)).mean(0)),
    f"risk difference per 100, clinic case mix (n={clinic.sum()})":
        q(100 * (risk_730(0, clinic) - risk_730(1, clinic)).mean(0)),
    f"risk difference per 100, low-risk patients (n={low_risk.sum()})":
        q(100 * (risk_730(0, low_risk) - risk_730(1, low_risk)).mean(0)),
}
ans_tab = pd.DataFrame(answers, index=["5%", "50%", "95%"]).T
ans_tab

fig, ax = plt.subplots(figsize=(8, 3.2))
labels_c = list(answers)[1:]
for i, lab in enumerate(labels_c):
    lo_, mid, hi_ = answers[lab]
    col = AQUA if "clinic" in lab else (RED if "misread" in lab else GREY)
    ax.errorbar(mid, i, xerr=[[mid - lo_], [hi_ - mid]], fmt="o", color=col, capsize=4, lw=2)
ax.set_yticks(range(len(labels_c)), [s.strip(" .") for s in labels_c], fontsize=9)
ax.invert_yaxis()
ax.set(xlabel="fewer events per 100 patients over 2 years (90% interval)",
       title="Same posterior, same trial - the population and the scale decide the number");