"""Phase 7 — publication figures (vector PDF) → paper/figures/. Palette and mark rules follow the dataviz reference
(categorical slots in fixed order, one hue sequential ramp, thin marks, direct labels, recessive grid, no dual axes)."""
import json, numpy as np, pandas as pd, matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.colors import LinearSegmentedColormap
from src.common import ROOT, path, get_logger
log = get_logger("figures"); P = path("processed"); F = ROOT / "paper" / "figures"; E = ROOT / "experiments"; F.mkdir(exist_ok=True)
C = ["#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#e87ba4", "#008300", "#6250d6", "#e34948"]
SEQ = LinearSegmentedColormap.from_list("seq", ["#cde2fb", "#86b6ef", "#3987e5", "#256abf", "#184f95", "#0d366b"])
TXT, TXT2, GRID = "#0b0b0b", "#52514e", "#e6e5e1"
plt.rcParams.update({"figure.constrained_layout.use": True, "font.size": 8.5, "axes.titlesize": 9, "axes.edgecolor": TXT2, "axes.labelcolor": TXT, "xtick.color": TXT2, "ytick.color": TXT2, "axes.spines.top": False, "axes.spines.right": False,
                     "axes.grid": True, "grid.color": GRID, "grid.linewidth": 0.6, "axes.axisbelow": True, "legend.frameon": False, "pdf.fonttype": 42})

def save(fig, name):
    fig.savefig(F / f"{name}.pdf", bbox_inches="tight"); plt.close(fig); log.info("wrote %s.pdf", name)

def fig_coverage():
    wk = pd.read_csv(P / "coverage_by_week.csv", parse_dates=["week"]); feed = pd.read_csv(P / "coverage_by_feed.csv"); feed = feed[feed.active > 0]
    fig, axes = plt.subplots(1, 2, figsize=(7.2, 3.0), gridspec_kw={"width_ratios": [1.3, 1]})
    ax = axes[0]; ax.plot(wk.week, wk.online, color=C[0], lw=2, marker="o", ms=4); ax.plot(wk.week, wk.circular, color=C[1], lw=2, marker="o", ms=4)
    ax.text(wk.week.iloc[-1], wk.online.iloc[-1], "  storefront", color=TXT, va="center"); ax.text(wk.week.iloc[-1], wk.circular.iloc[-1], "  circular", color=TXT, va="center")
    ax.set_ylabel("Listings observed in week"); ax.set_title("A. Weekly coverage by channel", loc="left", color=TXT); ax.tick_params(axis="x", rotation=30); ax.set_ylim(0); ax.xaxis.set_major_formatter(mdates.DateFormatter("%d %b")); ax.set_xlim(right=wk.week.iloc[-1] + pd.Timedelta(days=20))
    ax = axes[1]; comp = feed.groupby("company").active.sum().sort_values(ascending=True).tail(10)
    ax.barh(comp.index, comp.values, color=C[0], height=0.6); ax.set_title("B. Active listings by company (top 10)", loc="left", color=TXT); ax.set_xlabel("Active listings, 1 Nov 2026")
    for i, v in enumerate(comp.values): ax.text(v, i, f" {v:,}", va="center", color=TXT2, fontsize=8)
    ax.grid(axis="y", visible=False)
    save(fig, "fig_coverage")

def fig_attrition():
    att = pd.read_csv(P / "attrition_listings.csv"); v = pd.read_csv(P / "match_validation_summary.csv")
    fig, axes = plt.subplots(1, 2, figsize=(7.2, 3.0), gridspec_kw={"width_ratios": [1.4, 1]})
    ax = axes[0]; y = np.arange(len(att))[::-1]; ax.barh(y, att.listings, color=[C[0]] * (len(att) - 1) + [C[2]], height=0.6)
    labels = ["All listings", "excl. SAQ", "excl. pharmacies", "excl. non-food", "with a price", "excl. unrealistic", "parsable size", "matched (A–C)"]
    ax.set_yticks(y); ax.set_yticklabels(labels); ax.set_xlabel("Listings"); ax.set_title("A. Attrition flow", loc="left", color=TXT); ax.grid(axis="y", visible=False)
    for yi, n in zip(y, att.listings): ax.text(n, yi, f" {n:,}", va="center", color=TXT2, fontsize=8)
    ax = axes[1]; x = np.arange(4); ax.bar(x - 0.2, v.precision_strict, 0.38, color=C[0], label="Text review (240 pairs)"); ax.bar(x + 0.2, v.upc_precision.fillna(0), 0.38, color=C[1], label="UPC ground truth")
    ax.errorbar(x - 0.2, v.precision_strict, yerr=[np.maximum(0, v.precision_strict - v.ci_lo), np.maximum(0, v.ci_hi - v.precision_strict)], fmt="none", ecolor=TXT2, lw=1)
    ax.set_xticks(x); ax.set_xticklabels([f"Tier {t}" for t in v.tier]); ax.set_ylim(0, 1.05); ax.set_ylabel("Precision"); ax.set_title("B. Matching precision by tier", loc="left", color=TXT); ax.legend(loc="lower left", fontsize=7); ax.grid(axis="x", visible=False)
    save(fig, "fig_matching")

def fig_gaps():
    st = pd.read_csv(P / "modelA_groupweek_stats.csv"); fig, axes = plt.subplots(1, 2, figsize=(7.2, 3.0))
    ax = axes[0]; bins = np.linspace(0, 1.0, 41)
    ax.hist(st[st.same_company].gap.clip(upper=1), bins=bins, color=C[0], alpha=0.85, density=True, label="same company"); ax.hist(st[~st.same_company].gap.clip(upper=1), bins=bins, color=C[1], alpha=0.6, density=True, label="different companies")
    ax.set_xlabel("Within-product max–min log price gap (group-week)"); ax.set_ylabel("Density"); ax.legend(); ax.set_title("A. Distribution of cross-banner gaps", loc="left", color=TXT); ax.grid(axis="x", visible=False)
    fe = pd.read_csv(P / "modelA_banner_effects.csv").sort_values("coef"); fe = fe[fe.n_listing_weeks >= 50]
    ax = axes[1]; y = np.arange(len(fe)); ax.errorbar(fe.coef, y, xerr=1.96 * fe.se, fmt="o", color=C[0], ecolor=TXT2, ms=4, lw=1); ax.axvline(0, color=TXT2, lw=0.8)
    ax.set_yticks(y); ax.set_yticklabels(fe.banner); ax.set_xlabel("Banner effect on log price (ref. Metro), 95% CI"); ax.set_title("B. Relative price levels", loc="left", color=TXT); ax.grid(axis="y", visible=False)
    save(fig, "fig_dispersion")

def fig_heatmap():
    pt = pd.read_csv(P / "modelA_banner_pairs.csv"); pt = pt[pt.n >= 30]
    banners = sorted(set(pt.banner_r) | set(pt.banner_s)); M = pd.DataFrame(np.nan, index=banners, columns=banners)
    for _, r in pt.iterrows(): M.at[r.banner_r, r.banner_s] = r.median_gap; M.at[r.banner_s, r.banner_r] = -r.median_gap
    fig, ax = plt.subplots(figsize=(5.6, 4.8)); cm = LinearSegmentedColormap.from_list("div", ["#2a78d6", "#f0efec", "#e34948"])
    im = ax.imshow(M.values, cmap=cm, vmin=-0.3, vmax=0.3); ax.set_xticks(range(len(banners))); ax.set_xticklabels(banners, rotation=60, ha="right"); ax.set_yticks(range(len(banners))); ax.set_yticklabels(banners); ax.grid(False)
    for i in range(len(banners)):
        for j in range(len(banners)):
            if not np.isnan(M.values[i, j]): ax.text(j, i, f"{M.values[i, j]:+.2f}", ha="center", va="center", fontsize=6.5, color=TXT)
    cb = fig.colorbar(im, ax=ax, fraction=0.04); cb.set_label("Median log price gap, row banner minus column banner")
    ax.set_title("Median log gap, row minus column (pairs ≥ 30 group-weeks)", loc="left", color=TXT, fontsize=9)
    save(fig, "fig_heatmap")

def fig_promotions():
    es = pd.read_csv(P / "modelB_event_study.csv"); es["k"] = es.term.str.extract(r"(-?\d+)$").astype(int); es = pd.concat([es, pd.DataFrame([{"term": "ref", "coef": 0, "se": 0, "k": -1}])]).sort_values("k")
    ep = pd.read_parquet(P / "modelB_episodes.parquet"); tm = pd.read_csv(P / "modelB_transition_matrix.csv", index_col=0)
    fig, axes = plt.subplots(1, 3, figsize=(7.6, 2.8), gridspec_kw={"width_ratios": [1.3, 1, 0.8]})
    ax = axes[0]; ax.fill_between(es.k, es.coef - 1.96 * es.se, es.coef + 1.96 * es.se, color=C[0], alpha=0.18, lw=0); ax.plot(es.k, es.coef, color=C[0], lw=2, marker="o", ms=4); ax.axhline(0, color=TXT2, lw=0.8); ax.axvline(-0.5, color=TXT2, lw=0.8, ls=":")
    ax.set_xlabel("Weeks relative to first promoted week"); ax.set_ylabel("Log price vs week −1"); ax.set_title("A. Price path, temporary sales", loc="left", color=TXT)
    ax = axes[1]; unc = ep[~ep.censored]; ax.hist(unc[unc.kind == "return"].depth, bins=np.linspace(0, 1.2, 31), color=C[0], alpha=0.85, label="return (temporary)"); ax.hist(unc[unc.kind == "reset_down"].depth, bins=np.linspace(0, 1.2, 31), color=C[1], alpha=0.6, label="reset down")
    ax.set_xlabel("Episode depth, log(reference / price)"); ax.set_ylabel("Episodes"); ax.legend(fontsize=7); ax.set_title("B. Depth by episode kind", loc="left", color=TXT); ax.grid(axis="x", visible=False)
    ax = axes[2]; im = ax.imshow(tm.values, cmap=SEQ, vmin=0, vmax=1); ax.set_xticks([0, 1]); ax.set_xticklabels(["regular", "promo"]); ax.set_yticks([0, 1]); ax.set_yticklabels(["regular", "promo"]); ax.set_xlabel("state in week w"); ax.set_ylabel("state in week w−1"); ax.grid(False)
    for i in range(2):
        for j in range(2): ax.text(j, i, f"{tm.values[i, j]:.2f}", ha="center", va="center", color="white" if tm.values[i, j] > 0.5 else TXT)
    ax.set_title("C. Weekly transitions", loc="left", color=TXT)
    save(fig, "fig_promotions")

def fig_spillovers():
    d = pd.read_csv(P / "modelC_distributed_lags.csv"); fig, ax = plt.subplots(figsize=(5.2, 2.8))
    order = {"rival_all": 0, "rival_same": 1, "rival_cross": 2}; lab = {"rival_all": "all rivals", "rival_same": "same company", "rival_cross": "other company"}
    xs = ["_f2", "_f1", "", "_l1", "_l2", "_l3"]; xl = ["lead 2", "lead 1", "k=0", "lag 1", "lag 2", "lag 3"]
    for col, g in d.groupby("rival"):
        i = order[col]; vals = [g[g.term == col + s].coef.iloc[0] for s in xs]; ses = [g[g.term == col + s].se.iloc[0] for s in xs]
        ax.errorbar(np.arange(6) + (i - 1) * 0.22, vals, yerr=[1.96 * v for v in ses], fmt="o", color=C[i], ms=4, lw=1, label=lab[col])
    ax.axhline(0, color=TXT2, lw=0.8); ax.axvspan(-0.5, 1.5, color=GRID, alpha=0.5, lw=0); ax.set_xticks(range(6)); ax.set_xticklabels(xl); ax.set_ylabel("Response of own Δlog p"); ax.legend(fontsize=7, ncol=3, loc="upper left")
    ax.set_title("Own Δlog p on rivals' Δlog p (shaded = placebo leads)", loc="left", color=TXT, fontsize=9); ax.grid(axis="x", visible=False)
    save(fig, "fig_spillovers")

def fig_indices():
    d = pd.read_csv(P / "modelD_indices.csv", parse_dates=["week"]); s = d[d["sample"] == "storefront_posted"]
    fig, ax = plt.subplots(figsize=(5.6, 3.0))
    ax.plot(s.week, s.index_mean, color=C[1], lw=2, marker="o", ms=3); ax.plot(s.week, s["index"], color=C[0], lw=2, marker="o", ms=3); ax.plot(s.week, s.jevons, color=C[2], lw=2, marker="o", ms=3, ls="--"); ax.plot(s.week, s.carli, color=C[3], lw=1.5, marker="o", ms=3, ls=":")
    for col, name in [("index_mean", "naive mean of posted prices"), ("index", "chained Jevons (matched model)"), ("jevons", "fixed-basket Jevons"), ("carli", "fixed-basket Carli")]:
        ax.text(s.week.iloc[-1], s[col].iloc[-1], "  " + name, va="center", fontsize=7.5, color=TXT)
    ax.set_ylabel("Index (13 Aug 2026 = 100)"); ax.tick_params(axis="x", rotation=30); ax.xaxis.set_major_formatter(mdates.DateFormatter("%d %b")); ax.set_xlim(right=s.week.iloc[-1] + pd.Timedelta(days=45)); ax.set_title("Unweighted indices, storefront listings", loc="left", color=TXT, fontsize=9)
    save(fig, "fig_indices")

def fig_qgpai():
    q = pd.read_csv(P / "qgpai_components.csv", parse_dates=["week"]); fig, axes = plt.subplots(2, 2, figsize=(7.2, 4.8)); axes = axes.ravel()
    ax = axes[0]; ax.fill_between(q.week, q.cpi_lo, q.cpi_hi, color=C[0], alpha=0.18, lw=0); ax.plot(q.week, q.comparable_price_index, color=C[0], lw=2, marker="o", ms=3); ax.set_title("1. Comparable price index", loc="left", color=TXT); ax.set_ylabel("13 Aug = 100")
    ax = axes[1]; ax.plot(q.week, q.promo_share, color=C[1], lw=2, marker="o", ms=3); ax.plot(q.week, q.circular_depth_median, color=C[3], lw=2, marker="s", ms=3)
    ax.text(q.week.iloc[-1], q.promo_share.iloc[-1], "  share promoted\n  (storefront, inferred)", fontsize=7, va="center", color=TXT); ax.text(q.week.iloc[-1], q.circular_depth_median.iloc[-1], "  circular depth\n  (observed, median log)", fontsize=7, va="center", color=TXT)
    ax.set_title("2. Promotion intensity", loc="left", color=TXT); ax.set_ylim(0)
    ax = axes[2]; ax.fill_between(q.week, q.disp_lo, q.disp_hi, color=C[2], alpha=0.18, lw=0); ax.plot(q.week, q.dispersion_median_gap, color=C[2], lw=2, marker="o", ms=3); ax.set_title("3. Dispersion (median log gap)", loc="left", color=TXT); ax.set_ylim(0)
    ax = axes[3]
    for col, name, c, dy in [("basket_K1", "best single banner", C[0], 1.5), ("basket_K2", "two banners", C[1], 0.0), ("basket_unconstrained", "unconstrained", C[2], -1.5)]:
        ax.plot(q.week, q[col], color=c, lw=2, marker="o", ms=3); ax.text(q.week.iloc[-1], q[col].iloc[-1] + dy, "  " + name, fontsize=7, va="center", color=c)
    ax.set_title("4. Basket cost (CAD, fixed block)", loc="left", color=TXT); ax.set_ylim(0)
    for ax in axes: ax.tick_params(axis="x", rotation=30); ax.xaxis.set_major_formatter(mdates.DateFormatter("%d %b"))
    fig.suptitle("QGPAI components (component 5, access & affordability, not computable)", x=0.01, ha="left", color=TXT2, fontsize=9)
    save(fig, "fig_qgpai")

def fig_basket():
    r = pd.read_csv(P / "modelE_results.csv"); r5 = r[(r.B == 5) & r.single.isna() & r.fixed_cost.isna()]; fr = r5.groupby("K").cost.mean(); ks = [1, 2, 3, 4, 5]; costs = [fr.get(k, fr.get(99)) if k < 5 else fr.get(99) for k in ks]
    tc = pd.read_csv(P / "modelE_travel_sensitivity.csv")
    fig, axes = plt.subplots(1, 2, figsize=(7.2, 2.8))
    ax = axes[0]; ax.plot(ks, costs, color=C[0], lw=2, marker="o", ms=5); ax.set_xlabel("Maximum number of banners visited (K)"); ax.set_ylabel("Weekly basket cost (CAD, mean of 11 weeks)"); ax.set_xticks(ks)
    for k, c in zip(ks, costs): ax.text(k, c, f"  {c:.1f}", fontsize=7.5, color=TXT2, va="bottom")
    ax.set_title("A. Savings frontier (B = 5, 15 items)", loc="left", color=TXT)
    ax = axes[1]; ax.bar(tc.fixed_cost.astype(str), tc.cost_incl, color=C[1], width=0.55); ax.bar(tc.fixed_cost.astype(str), tc.cost_goods, color=C[0], width=0.55)
    for i, (g, t, n) in enumerate(zip(tc.cost_goods, tc.cost_incl, tc.n_used)): ax.text(i, t, f" {n:.1f} banners", ha="center", va="bottom", fontsize=7.5, color=TXT2)
    ax.set_xlabel("Illustrative fixed cost per banner visited (CAD)"); ax.set_ylabel("Goods (blue) + travel (orange)"); ax.set_title("B. Travel-cost sensitivity (assumption-based)", loc="left", color=TXT); ax.grid(axis="x", visible=False)
    save(fig, "fig_basket")

def fig_forecast():
    r = pd.read_csv(P / "modelF_rolling_results.csv"); fig, axes = plt.subplots(1, 2, figsize=(7.2, 2.8))
    ax = axes[0]; d = r[r.target == "log_price"]
    for i, (m, g) in enumerate(d.groupby("model")): ax.plot(g.origin_t, g.mae, color=C[i % 8], lw=1.8, marker="o", ms=3, label=m)
    ax.set_xlabel("Forecast origin (week index)"); ax.set_ylabel("MAE of next-week log price"); ax.legend(fontsize=6.5, ncol=2); ax.set_title("A. Price forecasts", loc="left", color=TXT)
    ax = axes[1]; d = r[r.target == "promotion"]
    for i, (m, g) in enumerate(d.groupby("model")): ax.plot(g.origin_t, g.auc, color=C[i % 8], lw=1.8, marker="o", ms=3, label=m)
    ax.set_xlabel("Forecast origin (week index)"); ax.set_ylabel("AUC, next-week promotion"); ax.set_ylim(0.4, 1); ax.legend(fontsize=6.5, ncol=2, loc="lower right"); ax.set_title("B. Promotion forecasts", loc="left", color=TXT)
    save(fig, "fig_forecast")

def fig_store_mixture():
    d = pd.read_parquet(P / "offers_daily.parquet", columns=["uid", "day", "p_min", "p_max", "p_mode", "n_distinct"]); x = d[d.uid == "metro:068400662600"].sort_values("day")
    pl = pd.read_parquet(path("interim") / "raw_price_log.parquet"); e = pl[pl.uid == "metro:068400662600"].copy(); e["ts"] = pd.to_datetime(e.ts, unit="s", utc=True).dt.tz_convert("America/Toronto")
    sm = pd.read_csv(P / "modelA_within_banner_mixture.csv").sort_values("share_multi", ascending=True).tail(12)
    fig, axes = plt.subplots(1, 2, figsize=(7.4, 2.9), gridspec_kw={"width_ratios": [1.4, 1]})
    ax = axes[0]; ax.step(e.ts, e.price, where="post", color=C[0], lw=1.2); ax.scatter(e.ts, e.price, color=C[0], s=8, zorder=3)
    ax.set_ylabel("Posted price (CAD)"); ax.set_title("A. One Metro listing: every logged price change", loc="left", color=TXT, fontsize=8.5); ax.tick_params(axis="x", rotation=30); ax.xaxis.set_major_formatter(mdates.DateFormatter("%d %b"))
    ax = axes[1]; ax.barh(sm.banner, sm.share_multi, color=C[0], height=0.6); ax.set_xlabel("Share of listing-weeks with > 1 price value"); ax.set_title("B. Multi-price weeks by banner", loc="left", color=TXT); ax.grid(axis="y", visible=False)
    for i, (v, r_) in enumerate(zip(sm.share_multi, sm.median_log_range_multi)): ax.text(v, i, f" {v:.2f} (range {r_:.2f})", va="center", fontsize=7, color=TXT2)
    save(fig, "fig_store_mixture")

def main():
    for f in [fig_coverage, fig_attrition, fig_gaps, fig_heatmap, fig_promotions, fig_spillovers, fig_indices, fig_qgpai, fig_basket, fig_forecast, fig_store_mixture]:
        try: f()
        except Exception as ex: log.error("%s failed: %s", f.__name__, ex)

if __name__ == "__main__":
    main()