"""Technical report analysis for grocery-sim's project site.

Reproducing this:

1. Generate the data with `grocery_sim`:

    from grocery_sim import GroceryStoreSimulation
    sim = GroceryStoreSimulation()
    sim.setup({
        "basic": {
            "name": "Technical Report Shop", "random_seed": 5501, "year": 3,
            "retain_earning": True, "retain_earning_from": "2026-01",
        },
        "events": {
            "war": ["2025-03-01", "2026-09-01"],
            "typhoon": "2025-07-15",
            "food_vat_cut": "2025-05-01",
            "tax_cut": "2026-02-01",
            "competitor": "2026-06-01",
            "operational_hazard": "2027-04-01",
        },
        "potential_investment": {
            "more_staff": True, "bigger_store": True, "upgrade_infrastructure": True,
        },
    })
    sim.simulate()
    data = sim.data()
    for name in data.keys():
        data[name].to_parquet(f"<DATA>/{name}.parquet")

2. Point DATA below at that folder, FIGDIR at wherever figures should land,
   and run this script. It writes every figure plus results.json — the
   single source every number in technical-report.qmd is transcribed from.

3. Dependencies beyond the package itself: statsmodels, scikit-learn,
   pymc, arviz, plotly, kaleido, openpyxl. On Windows, pm.sample() needs
   cores=1 (already set below) — the multiprocessing spawn backend
   otherwise re-imports this module in each worker and crashes without a
   `if __name__ == "__main__":` guard.

Runs GLM (NB regression + a difference-in-differences pass-through
design), ML (gradient-boosting demand forecast + a stock-out classifier),
and a hierarchical Bayesian partial-pooling model, against a real
three-year grocery_sim run. Every number quoted in the report prose is in
this script's own printed output / results.json -- nothing is hand-typed
from a notebook that could have drifted.
"""

import json
import os
import warnings

import numpy as np
import pandas as pd
import plotly.graph_objects as go

warnings.filterwarnings("ignore")

DATA = "./tech_report_data"       # <- point at the exported parquet folder (step 1 above)
FIGDIR = "./figures"              # <- where figures + results.json are written
os.makedirs(FIGDIR, exist_ok=True)

# --------------------------------------------------- palette (dataviz skill)
# validated: node scripts/validate_palette.js "#2a78d6,#e34948" --mode light
# -> ALL CHECKS PASS (lightness band, chroma floor, CVD sep 21.6, normal-
# vision floor 32.3, contrast). Fixed roles held across every chart in this
# report: blue = the model/processed estimate, red = the raw/unpooled/cost
# comparison, ink = actual observed data, muted = a de-emphasized baseline.
INK = "#0b0b0b"
SECONDARY = "#52514e"
MUTED = "#8c8c8c"
GRID = "#e5e5e3"
SURFACE = "#fcfcfb"
BLUE = "#2a78d6"       # slot 1 -- model / processed / pooled estimate
BLUE_WASH = "rgba(42,120,214,0.12)"
RED = "#e34948"        # slot 8 -- raw / unpooled / cost comparison

PLOT = dict(
    plot_bgcolor=SURFACE,
    paper_bgcolor=SURFACE,
    font=dict(color=INK, size=13, family="Helvetica, Arial, sans-serif"),
    margin=dict(l=64, r=28, t=56, b=52),
)
AXIS = dict(
    showgrid=True, gridcolor=GRID, gridwidth=1, zeroline=False,
    showline=True, linecolor=GRID, ticks="outside", tickcolor=GRID,
    tickfont=dict(color=SECONDARY),
)


def style(fig, title=None, height=420, width=920, showlegend=True):
    fig.update_layout(
        **PLOT, height=height, width=width, showlegend=showlegend,
        legend=dict(bgcolor="rgba(0,0,0,0)", font=dict(color=SECONDARY, size=12)),
        title=dict(text=title, x=0, xanchor="left", font=dict(color=INK, size=15)) if title else None,
    )
    fig.update_xaxes(**AXIS, title_font=dict(color=SECONDARY))
    fig.update_yaxes(**AXIS, title_font=dict(color=SECONDARY))
    return fig


def savefig(fig, name, height=420, width=920, showlegend=True, title=None):
    style(fig, title=title, height=height, width=width, showlegend=showlegend)
    fig.write_image(f"{FIGDIR}/{name}.png", scale=2)


results = {}

# ==================================================================== load
receipts = pd.read_parquet(f"{DATA}/receipts.parquet")
weather = pd.read_parquet(f"{DATA}/weather.parquet")
calendar = pd.read_parquet(f"{DATA}/calendar.parquet")
cost_sheet = pd.read_parquet(f"{DATA}/cost_sheet.parquet")
procurement = pd.read_parquet(f"{DATA}/procurement.parquet")
price_history = pd.read_parquet(f"{DATA}/price_history.parquet")
inventory_eod = pd.read_parquet(f"{DATA}/inventory_eod.parquet")
skus = pd.read_excel("SKUs.xlsx")  # ships inside package/grocery_sim/SKUs.xlsx

uid_cat = skus.set_index("uid")["category"]
receipts["category"] = receipts["uid"].map(uid_cat)
procurement["category"] = procurement["uid"].map(uid_cat)
price_history["category"] = price_history["uid"].map(uid_cat)
inventory_eod["category"] = inventory_eod["uid"].map(uid_cat) if "uid" in inventory_eod.columns else None

for df in (receipts, weather, calendar, procurement, price_history, inventory_eod):
    if "date" in df.columns:
        df["date"] = pd.to_datetime(df["date"])

print("loaded:", receipts.shape, weather.shape, calendar.shape, cost_sheet.shape)

# ============================================================ GLM section
import statsmodels.api as sm
import statsmodels.formula.api as smf

sales = receipts[receipts["qty"] > 0].copy()
daily_units = sales.groupby("date")["qty"].sum().rename("units").to_frame()
daily = daily_units.join(weather.set_index("date")).join(calendar.set_index("date"))
daily = daily[daily["closed"] == 0].copy()
daily["t"] = (daily.index - daily.index.min()).days
daily["dow"] = daily.index.dayofweek
daily["is_weekend"] = (daily["dow"] >= 5).astype(int)
daily["is_pre_holiday"] = daily["pre_holiday"].astype(int)
temp_seasonal = daily.groupby(daily.index.dayofyear)["temp_C"].transform("mean")
daily["temp_anom"] = daily["temp_C"] - temp_seasonal

glm_df = daily.dropna(subset=["units", "temp_anom", "rain_mm", "wet"]).copy()

nb_model = smf.glm(
    formula="units ~ temp_anom + rain_mm + wet + is_weekend + is_pre_holiday + t",
    data=glm_df,
    family=sm.families.NegativeBinomial(),
).fit(cov_type="HAC", cov_kwds={"maxlags": 7})

glm_table = pd.DataFrame({
    "coef": nb_model.params,
    "IRR": np.exp(nb_model.params),
    "se_HAC": nb_model.bse,
    "p_value": nb_model.pvalues,
    "ci_low": np.exp(nb_model.conf_int()[0]),
    "ci_high": np.exp(nb_model.conf_int()[1]),
})
print(glm_table.round(4))
results["glm_demand"] = {
    "n_obs": int(nb_model.nobs),
    "pseudo_r2": float(1 - nb_model.deviance / nb_model.null_deviance),
    "table": glm_table.round(4).to_dict(orient="index"),
}

# fitted vs actual plot (7-day rolling mean, so weekly cycle doesn't dominate)
glm_df["fitted"] = nb_model.fittedvalues
roll = glm_df[["units", "fitted"]].rolling(7, min_periods=1).mean()
fig = go.Figure()
fig.add_trace(go.Scatter(
    x=roll.index, y=roll["units"], name="actual (7-day mean)",
    mode="lines", line=dict(color=INK, width=2),
))
fig.add_trace(go.Scatter(
    x=roll.index, y=roll["fitted"], name="GLM fitted",
    mode="lines", line=dict(color=BLUE, width=2),
))
fig.update_yaxes(title="units/day")
savefig(fig, "01_glm_fitted_vs_actual", title="Negative-binomial demand GLM: fitted vs. actual")

# temperature partial-effect plot
temp_grid = np.linspace(glm_df["temp_anom"].min(), glm_df["temp_anom"].max(), 60)
base = glm_df[["rain_mm", "wet", "is_weekend", "is_pre_holiday", "t"]].mean()
pred_df = pd.DataFrame({"temp_anom": temp_grid, **{k: base[k] for k in base.index}})
pred = nb_model.get_prediction(pred_df).summary_frame()
fig = go.Figure()
fig.add_trace(go.Scatter(
    x=np.concatenate([temp_grid, temp_grid[::-1]]),
    y=np.concatenate([pred["mean_ci_upper"], pred["mean_ci_lower"][::-1]]),
    fill="toself", fillcolor=BLUE_WASH, line=dict(width=0),
    showlegend=False, hoverinfo="skip",
))
fig.add_trace(go.Scatter(
    x=temp_grid, y=pred["mean"], mode="lines", line=dict(color=BLUE, width=2),
    showlegend=False,
))
fig.update_xaxes(title="temperature anomaly vs. day-of-year normal (°C)")
fig.update_yaxes(title="predicted units/day")
savefig(fig, "02_glm_temp_partial_effect",
        title="Partial effect of temperature anomaly on demand (other covariates at their mean)",
        showlegend=False)

# ---------------------------------------------- pass-through: DiD design
# events.food_vat_cut halves VAT on the reduced-rate (food) group on one
# exact date -- but a naive before/after event study on a food category
# alone is confounded here by the concurrent events.war cost shock
# (started 2025-03-01, ~120-day decay, hits every category), which by
# itself pushes the realized cost step to ~-8% against a theoretical
# -4.55%. A standard-VAT category (unaffected by the VAT cut, but equally
# exposed to the war shock) lets that common trend be differenced out --
# a difference-in-differences design, not just a cleaner window.
treated_cat = "Pantry Staples and Packaged Goods"   # reduced VAT -> cut
control_cat = "Household and Cleaning Supplies"      # standard VAT -> unaffected
vat_cut_date = pd.Timestamp("2025-05-01")


all_weeks = pd.period_range(price_history["date"].min(), price_history["date"].max(), freq="W").start_time


def weekly_price_level(cat):
    """A proper held-price index: price_history only logs a row when a
    price *changes* (menu-cost hysteresis means most SKUs reprice only a
    few times a year), so naively averaging whatever rows exist in a
    given week means some weeks are an average of 1-2 SKUs that happened
    to reprice, not the category's actual shelf level that week. Forward-
    fill each SKU's last known price onto every week instead, so the
    weekly average is always across the category's full SKU set."""
    ph_cat = price_history[price_history["category"] == cat].copy()
    ph_cat["week"] = ph_cat["date"].dt.to_period("W").dt.start_time
    panel = ph_cat.pivot_table(index="week", columns="uid", values="price", aggfunc="last")
    panel = panel.reindex(all_weeks).ffill()
    return panel.mean(axis=1).rename("price")


def weekly_cost_index(cat):
    proc_cat = procurement[procurement["category"] == cat].copy()
    proc_cat["week"] = pd.to_datetime(proc_cat["delivery_date"]).dt.to_period("W").dt.start_time
    return proc_cat.groupby("week")["unit_cost"].mean().reindex(all_weeks).ffill().rename("cost")


def weekly_index(cat):
    df = pd.concat([weekly_price_level(cat), weekly_cost_index(cat)], axis=1).dropna()
    df["weeks_since"] = ((df.index - vat_cut_date).days // 7)
    return df[(df["weeks_since"] >= -8) & (df["weeks_since"] <= 8)].copy()


treated = weekly_index(treated_cat)
control = weekly_index(control_cat)


def pre_post_means(df):
    pre = df[df["weeks_since"] < 0][["price", "cost"]].mean()
    post = df[(df["weeks_since"] >= 0) & (df["weeks_since"] <= 7)][["price", "cost"]].mean()
    return pre, post


t_pre, t_post = pre_post_means(treated)
c_pre, c_post = pre_post_means(control)

d_cost_treated = np.log(t_post["cost"]) - np.log(t_pre["cost"])
d_cost_control = np.log(c_post["cost"]) - np.log(c_pre["cost"])
did_cost = d_cost_treated - d_cost_control  # the VAT-cut-specific cost move

d_price_treated = np.log(t_post["price"]) - np.log(t_pre["price"])
d_price_control = np.log(c_post["price"]) - np.log(c_pre["price"])
did_price = d_price_treated - d_price_control  # the VAT-cut-specific price move

did_pass_through = did_price / did_cost

print("naive treated-only cost drop:", np.exp(d_cost_treated) - 1, "(confounded by the war shock)")
print("control-category cost move (the common trend):", np.exp(d_cost_control) - 1)
print("DiD cost move (VAT-cut-specific):", np.exp(did_cost) - 1, "vs theoretical", 1.05 / 1.10 - 1)
print("DiD price move (VAT-cut-specific):", np.exp(did_price) - 1)
print("DiD pass-through (uninformative -- see report):", did_pass_through)

results["pass_through"] = {
    "treated_category": treated_cat,
    "control_category": control_cat,
    "theoretical_cost_drop": round(1.05 / 1.10 - 1, 4),
    "naive_treated_cost_drop": round(float(np.exp(d_cost_treated) - 1), 4),
    "control_cost_move": round(float(np.exp(d_cost_control) - 1), 4),
    "did_cost_drop": round(float(np.exp(did_cost) - 1), 4),
    "did_price_drop": round(float(np.exp(did_price) - 1), 4),
    "did_pass_through": round(float(did_pass_through), 3),
}

for df in (treated, control):
    df["cost_idx100"] = 100 * np.exp(np.log(df["cost"]) - np.log(df.loc[df["weeks_since"] < 0, "cost"]).mean())
    df["price_idx100"] = 100 * np.exp(np.log(df["price"]) - np.log(df.loc[df["weeks_since"] < 0, "price"]).mean())

from plotly.subplots import make_subplots

fig = make_subplots(
    rows=1, cols=2, shared_yaxes=True,
    subplot_titles=(f"{treated_cat} (treated: VAT cut)", f"{control_cat} (control: no VAT change)"),
    horizontal_spacing=0.06,
)
for col, df in enumerate((treated, control), start=1):
    show = col == 1
    fig.add_trace(go.Scatter(
        x=df["weeks_since"], y=df["cost_idx100"], name="invoice cost",
        mode="lines+markers", line=dict(color=RED, width=2), marker=dict(size=8, line=dict(color=SURFACE, width=2)),
        legendgroup="cost", showlegend=show,
    ), row=1, col=col)
    fig.add_trace(go.Scatter(
        x=df["weeks_since"], y=df["price_idx100"], name="shelf price",
        mode="lines+markers", line=dict(color=BLUE, width=2), marker=dict(size=8, line=dict(color=SURFACE, width=2)),
        legendgroup="price", showlegend=show,
    ), row=1, col=col)
    fig.add_vline(x=0, line=dict(color=MUTED, width=1, dash="dash"), row=1, col=col)
    fig.update_xaxes(title="weeks since the food-VAT cut", row=1, col=col, **AXIS)
fig.update_yaxes(title="index (pre-event mean = 100)", row=1, col=1, **AXIS)
fig.update_yaxes(row=1, col=2, **AXIS)
for ann in fig.layout.annotations:
    ann.font = dict(color=INK, size=13)
fig.update_layout(**PLOT, height=440, width=980,
                   legend=dict(bgcolor="rgba(0,0,0,0)", font=dict(color=SECONDARY, size=12), x=0.01, y=0.02),
                   title=dict(text="Cost pass-through, difference-in-differences design", x=0, xanchor="left", font=dict(color=INK, size=15)))
fig.write_image(f"{FIGDIR}/03_pass_through_did.png", scale=2)

print("GLM section done")

# ===================================================================== ML
from sklearn.ensemble import HistGradientBoostingRegressor, HistGradientBoostingClassifier
from sklearn.metrics import mean_absolute_error, mean_squared_error, average_precision_score, precision_recall_curve

forecast_cat = "Dairy and Eggs"
cat_sales = sales[sales["category"] == forecast_cat].copy()
weekly = cat_sales.groupby(pd.Grouper(key="date", freq="W"))["qty"].sum().rename("units").to_frame()
weekly = weekly.join(weather.set_index("date").resample("W").mean()[["temp_C", "rain_mm"]])
cal_w = calendar.set_index("date")
weekly["holiday_week"] = cal_w["holiday"].notna().resample("W").sum().reindex(weekly.index, fill_value=0)
weekly["week_of_year"] = weekly.index.isocalendar().week.astype(int)
weekly["month"] = weekly.index.month
for lag in (1, 2, 52):
    weekly[f"lag_{lag}"] = weekly["units"].shift(lag)
weekly["roll4"] = weekly["units"].shift(1).rolling(4).mean()
weekly = weekly.dropna()

train = weekly[weekly.index.year < 2027]
test = weekly[weekly.index.year == 2027]
features = ["temp_C", "rain_mm", "holiday_week", "week_of_year", "month", "lag_1", "lag_2", "lag_52", "roll4"]

gbm = HistGradientBoostingRegressor(max_depth=3, learning_rate=0.08, max_iter=200, random_state=0)
gbm.fit(train[features], train["units"])
pred_gbm = gbm.predict(test[features])

naive_pred = test["lag_52"].values  # seasonal-naive: same week last year

mae_gbm = mean_absolute_error(test["units"], pred_gbm)
mae_naive = mean_absolute_error(test["units"], naive_pred)
rmse_gbm = mean_squared_error(test["units"], pred_gbm) ** 0.5
rmse_naive = mean_squared_error(test["units"], naive_pred) ** 0.5

print(f"GBM  MAE={mae_gbm:.1f} RMSE={rmse_gbm:.1f}")
print(f"naive MAE={mae_naive:.1f} RMSE={rmse_naive:.1f}")
print("MAE improvement over naive:", 1 - mae_gbm / mae_naive)

results["ml_forecast"] = {
    "category": forecast_cat,
    "n_train": int(len(train)),
    "n_test": int(len(test)),
    "mae_gbm": round(float(mae_gbm), 1),
    "mae_naive": round(float(mae_naive), 1),
    "rmse_gbm": round(float(rmse_gbm), 1),
    "rmse_naive": round(float(rmse_naive), 1),
    "mae_improvement_pct": round(float(1 - mae_gbm / mae_naive) * 100, 1),
    "feature_importance": {},
}

fig = go.Figure()
fig.add_trace(go.Scatter(
    x=test.index, y=test["units"], name="actual", mode="lines+markers",
    line=dict(color=INK, width=2), marker=dict(size=7, line=dict(color=SURFACE, width=1.5)),
))
fig.add_trace(go.Scatter(
    x=test.index, y=pred_gbm, name="gradient boosting", mode="lines+markers",
    line=dict(color=BLUE, width=2), marker=dict(size=7, line=dict(color=SURFACE, width=1.5)),
))
fig.add_trace(go.Scatter(
    x=test.index, y=naive_pred, name="seasonal-naive (52-wk lag)", mode="lines",
    line=dict(color=MUTED, width=2, dash="dash"),
))
fig.update_yaxes(title="units/week")
savefig(fig, "04_ml_forecast_holdout", title=f"2027 holdout forecast — {forecast_cat}")

# permutation-free importance from the trained HGB via built-in feature contribution proxy:
# HistGradientBoostingRegressor has no feature_importances_, so use permutation importance
from sklearn.inspection import permutation_importance
perm = permutation_importance(gbm, test[features], test["units"], n_repeats=30, random_state=0)
imp = pd.Series(perm.importances_mean, index=features).sort_values()
results["ml_forecast"]["feature_importance"] = imp.round(3).to_dict()
print(imp)

fig = go.Figure(go.Bar(
    x=imp.values, y=imp.index, orientation="h",
    marker=dict(color=BLUE),
))
fig.update_xaxes(title="permutation importance (MAE increase when shuffled)")
savefig(fig, "05_ml_feature_importance", title="What the forecaster actually relies on",
        height=440, showlegend=False)

print("ML forecast section done")

# ---------------------------------------------------- stockout classifier
stockout_days = inventory_eod[inventory_eod["on_hand"] == 0][["uid", "date"]].copy()
stockout_days["week"] = stockout_days["date"].dt.to_period("W").dt.start_time
stockout_weeks = stockout_days.groupby(["uid", "week"]).size().reset_index(name="n")
stockout_weeks["stockout"] = 1

inv = inventory_eod.copy()
inv["week"] = inv["date"].dt.to_period("W").dt.start_time
weekly_inv = inv.groupby(["uid", "week"])["on_hand"].mean().rename("avg_on_hand").reset_index()

sales_by_sku_week = sales.copy()
sales_by_sku_week["week"] = sales_by_sku_week["date"].dt.to_period("W").dt.start_time
weekly_demand = sales_by_sku_week.groupby(["uid", "week"])["qty"].sum().rename("weekly_units").reset_index()

panel = weekly_inv.merge(weekly_demand, on=["uid", "week"], how="left").fillna({"weekly_units": 0})
panel["category"] = panel["uid"].map(uid_cat)
panel = panel.sort_values(["uid", "week"])
panel["avg_daily_demand_lag"] = panel.groupby("uid")["weekly_units"].shift(1) / 7.0
panel["cover_days"] = panel["avg_on_hand"] / panel["avg_daily_demand_lag"].replace(0, np.nan)
panel["next_week"] = panel.groupby("uid")["week"].shift(-1)
panel = panel.merge(
    stockout_weeks[["uid", "week", "stockout"]].rename(columns={"week": "next_week"}),
    on=["uid", "next_week"], how="left",
)
panel["stockout"] = panel["stockout"].fillna(0).astype(int)
panel["month"] = panel["week"].dt.month
panel = panel.dropna(subset=["cover_days"])
panel = panel[np.isfinite(panel["cover_days"])]
panel["cat_code"] = panel["category"].astype("category").cat.codes

cls_features = ["cover_days", "avg_daily_demand_lag", "month", "cat_code"]
panel_train = panel[panel["week"].dt.year < 2027]
panel_test = panel[panel["week"].dt.year == 2027]

clf = HistGradientBoostingClassifier(max_depth=4, learning_rate=0.08, max_iter=150, random_state=0)
clf.fit(panel_train[cls_features], panel_train["stockout"])
proba = clf.predict_proba(panel_test[cls_features])[:, 1]

pr_auc = average_precision_score(panel_test["stockout"], proba)
base_rate = panel_test["stockout"].mean()
precision, recall, thresh = precision_recall_curve(panel_test["stockout"], proba)

print("stockout base rate (test):", base_rate, "PR-AUC:", pr_auc)
results["ml_stockout"] = {
    "n_train": int(len(panel_train)),
    "n_test": int(len(panel_test)),
    "base_rate_test": round(float(base_rate), 4),
    "pr_auc": round(float(pr_auc), 4),
}

fig = go.Figure()
fig.add_trace(go.Scatter(
    x=recall, y=precision, name="classifier", mode="lines", line=dict(color=BLUE, width=2),
))
fig.add_trace(go.Scatter(
    x=[0, 1], y=[base_rate, base_rate], name=f"random baseline (base rate = {base_rate:.3f})",
    mode="lines", line=dict(color=MUTED, width=2, dash="dash"),
))
fig.update_xaxes(title="recall", range=[0, 1])
fig.update_yaxes(title="precision", range=[0, 1.02])
savefig(fig, "06_ml_stockout_pr_curve", title=f"Stockout-risk classifier: PR curve (AUC = {pr_auc:.3f})",
        height=460, width=640)

print("ML section done")

# ============================================================ Bayesian
# hierarchical partial pooling of the weekend demand effect, by category.
# Deliberately restricted to the first 90 trading days (not the full three
# years): with three years of data every category's own OLS estimate is
# already precise enough that pooling barely changes anything, which would
# make a weak demonstration. A new shop's first quarter is exactly when a
# real analyst faces this problem for real -- some categories have enough
# volume to trust on their own, others don't yet, and partial pooling is
# the honest way to use what the whole shop knows to steady the ones that
# don't.
early = sales[sales["date"] < sales["date"].min() + pd.Timedelta(days=90)].copy()
early_cal = calendar.set_index("date")
daily_cat = early.groupby(["category", "date"])["qty"].sum().rename("units").reset_index()
daily_cat["is_weekend"] = daily_cat["date"].map(early_cal["dow"]).ge(5).astype(int)
daily_cat["log_units"] = np.log1p(daily_cat["units"])

cats = sorted(daily_cat["category"].unique())
unpooled = []
for c in cats:
    sub = daily_cat[daily_cat["category"] == c]
    m = smf.ols("log_units ~ is_weekend", data=sub).fit()
    unpooled.append({
        "category": c,
        "n_days": int(len(sub)),
        "beta": float(m.params["is_weekend"]),
        "se": float(m.bse["is_weekend"]),
    })
unpooled = pd.DataFrame(unpooled).sort_values("se").reset_index(drop=True)
print(unpooled)

import pymc as pm

with pm.Model() as hier:
    mu = pm.Normal("mu", 0, 1)
    tau = pm.HalfNormal("tau", 0.5)
    beta_true = pm.Normal("beta_true", mu, tau, shape=len(unpooled))
    pm.Normal("obs", beta_true, unpooled["se"].values, observed=unpooled["beta"].values)
    # cores=1: avoids the Windows multiprocessing spawn crash (see module
    # docstring) -- fine here, the model is tiny (12 observations)
    idata = pm.sample(2000, tune=1500, chains=4, cores=1, target_accept=0.95, random_seed=0, progressbar=False)

post_mean = idata.posterior["beta_true"].mean(dim=("chain", "draw")).values
post_lo = idata.posterior["beta_true"].quantile(0.025, dim=("chain", "draw")).values
post_hi = idata.posterior["beta_true"].quantile(0.975, dim=("chain", "draw")).values
global_mu = float(idata.posterior["mu"].mean())
rhat_max = float(pm.summary(idata, var_names=["mu", "tau"])["r_hat"].max())

unpooled["pooled_mean"] = post_mean
unpooled["pooled_lo"] = post_lo
unpooled["pooled_hi"] = post_hi
print(unpooled)
print("global mu:", global_mu, "max r_hat:", rhat_max)

results["bayes_hierarchical"] = {
    "n_days_used": 90,
    "n_categories": len(cats),
    "global_weekend_effect_log": round(global_mu, 3),
    "global_weekend_effect_pct": round((np.exp(global_mu) - 1) * 100, 1),
    "max_rhat": round(rhat_max, 4),
    "table": unpooled.round(3).to_dict(orient="records"),
}

y_labels = [f"{c}  (n={n})" for c, n in zip(unpooled["category"], unpooled["n_days"])]
y = np.arange(len(unpooled))
fig = go.Figure()
for i in range(len(unpooled)):
    fig.add_trace(go.Scatter(
        x=[unpooled["pooled_lo"].iloc[i], unpooled["pooled_hi"].iloc[i]], y=[y[i] - 0.15, y[i] - 0.15],
        mode="lines", line=dict(color=BLUE, width=6), opacity=0.25, showlegend=False, hoverinfo="skip",
    ))
    fig.add_trace(go.Scatter(
        x=[unpooled["beta"].iloc[i], unpooled["pooled_mean"].iloc[i]], y=[y[i] + 0.15, y[i] - 0.15],
        mode="lines", line=dict(color=MUTED, width=1), showlegend=False, hoverinfo="skip",
    ))
fig.add_trace(go.Scatter(
    x=unpooled["beta"], y=y + 0.15, name="unpooled (per-category OLS)", mode="markers",
    marker=dict(color=RED, size=10, line=dict(color=SURFACE, width=2)),
))
fig.add_trace(go.Scatter(
    x=unpooled["pooled_mean"], y=y - 0.15, name="partial-pooled (hierarchical posterior mean)", mode="markers",
    marker=dict(color=BLUE, size=10, line=dict(color=SURFACE, width=2)),
))
fig.add_vline(x=global_mu, line=dict(color=INK, width=1, dash="dash"))
fig.update_yaxes(tickmode="array", tickvals=list(y), ticktext=y_labels, title=None)
fig.update_xaxes(title="weekend effect on log(units) — first 90 days")
plot_kwargs = dict(PLOT)
plot_kwargs["margin"] = dict(l=230, r=28, t=80, b=52)
fig.update_layout(**plot_kwargs, height=560, width=980,
                   legend=dict(bgcolor="rgba(0,0,0,0)", font=dict(color=SECONDARY, size=12),
                               orientation="h", y=1.1, x=0),
                   title=dict(text="Partial pooling shrinks noisy category estimates toward the shop-wide mean",
                              x=0, xanchor="left", font=dict(color=INK, size=14)))
fig.update_xaxes(**AXIS)
fig.update_yaxes(**AXIS, tickmode="array", tickvals=list(y), ticktext=y_labels)
fig.write_image(f"{FIGDIR}/07_bayes_shrinkage.png", scale=2)

print("Bayesian section done")

with open(f"{FIGDIR}/results.json", "w") as f:
    json.dump(results, f, indent=2, default=str)
print("wrote results.json")