"""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
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({"font.size": 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 = 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=[v.precision_strict - v.ci_lo, 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("Pairwise median gaps on identical products (pairs with ≥ 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 around 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("Lead–lag association with rivals' price changes (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.set_title("Unweighted price 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 (matched groups)", 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. Retail price dispersion (median log gap)", loc="left", color=TXT); ax.set_ylim(0)
    ax = axes[3]
    for col, name, c in [("basket_K1", "best single banner", C[0]), ("basket_K2", "two banners", C[1]), ("basket_unconstrained", "unconstrained", C[2])]:
        ax.plot(q.week, q[col], color=c, lw=2, marker="o", ms=3); ax.text(q.week.iloc[-1], q[col].iloc[-1], "  " + name, fontsize=7, va="center", color=TXT)
    ax.set_title("4. Feasible basket cost (CAD, fixed block)", loc="left", color=TXT); ax.set_ylim(0)
    for ax in axes: ax.tick_params(axis="x", rotation=30)
    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 listing (Metro, mayonnaise 890 mL): every logged price change", loc="left", color=TXT, fontsize=8.5); ax.tick_params(axis="x", rotation=30)
    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()