"""Phase 5 — Quebec Grocery Price and Affordability Index (QGPAI): components published separately (preregistration §3).

1 Comparable Price Index  — chained Jevons over matched groups (A+B), storefront, base week 2026-08-13 = 100
2 Promotion Intensity     — share of listing-weeks promoted (inferred, storefront) and observed circular depth
3 Retail Price Dispersion — weekly median within-group max–min log gap
4 Feasible Basket Cost    — Model E weekly cost at K=1 (best single banner) and K=2 on a block that is constant over weeks
                            (5 widest-coverage banners × items they all carry in every week)
5 Access & Affordability  — NOT computed (no geography / income join); documented placeholder
Illustrative composite: equal-weight z-scores of components 1–4 (and a 0.5/0.5 alternative on 1 and 4) — labelled as an
illustration, with listing-level bootstrap bands for components 1 and 3.
"""
import json, numpy as np, pandas as pd
from src.common import path, ROOT, get_logger, ledger, seed_everything, SEED
log = get_logger("qgpai"); P = path("processed"); E = ROOT / "experiments"

def jevons(df, col="p_mode"):
    weeks = sorted(df.week.unique()); piv = df.pivot_table(index="uid", columns="week", values=col); idx = [100.0]
    for a, b in zip(weeks, weeks[1:]):
        both = piv[[a, b]].dropna(); idx.append(idx[-1] * np.exp((np.log(both[b]) - np.log(both[a])).mean()))
    return pd.Series(idx, index=weeks)

def main():
    seed_everything(); rng = np.random.default_rng(SEED); out = {}
    pm = pd.read_parquet(P / "panel_matched_AB.parquet"); pm = pm[(pm.channel == "online_storefront") & (~pm.ambiguous)]
    pm = pm[pm.groupby(["group_id", "week"]).banner.transform("nunique") >= 2]
    # 1 comparable price index (listings inside matched groups) with bootstrap over groups
    c1 = jevons(pm); groups = pm.group_id.unique(); boots = []
    for b in range(299):
        s = rng.choice(groups, size=len(groups), replace=True); d = pd.concat([pm[pm.group_id == g] for g in s[:0]]) if False else None
        w = pd.Series(rng.multinomial(len(groups), np.full(len(groups), 1 / len(groups))), index=groups)
        # weighted Jevons: weight listings by their group's multiplicity
        piv = pm.pivot_table(index="uid", columns="week", values="p_mode"); gw = pm.drop_duplicates("uid").set_index("uid").group_id.map(w).reindex(piv.index).values
        weeks = list(piv.columns); idx = [100.0]
        for a, c in zip(weeks, weeks[1:]):
            both = piv[[a, c]].dropna(); ww = gw[piv.index.get_indexer(both.index)]
            idx.append(idx[-1] * np.exp(np.average(np.log(both[c]) - np.log(both[a]), weights=ww)))
        boots.append(idx)
    boots = np.array(boots); c1_lo, c1_hi = np.percentile(boots, 2.5, axis=0), np.percentile(boots, 97.5, axis=0)
    # 3 dispersion with bootstrap over groups
    gaps = pm.groupby(["group_id", "week"]).log_p.agg(lambda s: s.max() - s.min()).reset_index()
    c3 = gaps.groupby("week").log_p.median(); b3 = []
    for b in range(299):
        w = pd.Series(rng.multinomial(len(groups), np.full(len(groups), 1 / len(groups))), index=groups)
        g2 = gaps.loc[gaps.index.repeat(gaps.group_id.map(w).values)]; b3.append(g2.groupby("week").log_p.median().reindex(c3.index).values)
    b3 = np.array(b3); c3_lo, c3_hi = np.percentile(b3, 2.5, axis=0), np.percentile(b3, 97.5, axis=0)
    # 2 promotion intensity (storefront inferred) — recompute weekly from panel_weekly with the same 8-week-mode rule
    pw = pd.read_parquet(P / "panel_weekly.parquet"); sf = pw[pw.channel == "online_storefront"].sort_values(["uid", "week"]).copy()
    sf = sf[sf.groupby("uid").week.transform("nunique") >= 8]
    def rolling_mode(s):
        v = s.tolist(); o = []
        for i in range(len(v)):
            vc = pd.Series(v[max(0, i - 7): i + 1]).value_counts(); o.append(max(vc[vc == vc.max()].index))
        return pd.Series(o, index=s.index)
    sf["ref"] = sf.groupby("uid").p_mode.transform(rolling_mode); sf["promo"] = sf.p_mode <= 0.9 * sf.ref; sf["depth"] = np.where(sf.promo, np.log(sf.ref / sf.p_mode), np.nan)
    c2 = sf.groupby("week").agg(promo_share=("promo", "mean"), depth_mean=("depth", "mean"), depth_sd=("depth", "std"))
    c2.iloc[0, :] = np.nan   # first week has no trailing history: promotion state undefined
    # circular observed depth by week (items valid that week)
    prod = pd.read_parquet(P / "products.parquet"); ret = pd.read_parquet(P / "retailers.parquet")
    circ = prod.merge(ret[["feed_id", "channel", "is_grocery"]], left_on="source", right_on="feed_id"); circ = circ[circ.channel.str.startswith("circular") & circ.is_grocery & circ.regular_price.notna() & circ.price.notna()]
    circ["vf"] = pd.to_datetime(circ.valid_from, utc=True, errors="coerce").dt.tz_convert("America/Toronto").dt.tz_localize(None).dt.normalize()
    circ["depth"] = np.log(circ.regular_price / circ.price); cd = circ.groupby("vf").depth.agg(["median", "size"]); cd.index.name = "week"
    # 4 basket cost
    e = pd.read_csv(P / "modelE_fixed_block.csv"); e["week"] = pd.to_datetime(e.week)   # constant item set over weeks
    c4 = e.pivot_table(index="week", columns="K", values="cost").rename(columns={1: "basket_K1", 2: "basket_K2", 99: "basket_unconstrained"})
    # assemble
    weeks = pd.to_datetime(pd.Series(c1.index))
    q = pd.DataFrame({"week": weeks, "comparable_price_index": c1.values, "cpi_lo": c1_lo, "cpi_hi": c1_hi, "dispersion_median_gap": c3.reindex(c1.index).values, "disp_lo": c3_lo, "disp_hi": c3_hi})
    q = q.merge(c2.reset_index().assign(week=lambda d: pd.to_datetime(d.week)), on="week", how="left").merge(c4.reset_index(), on="week", how="left")
    q = q.merge(cd.reset_index().rename(columns={"median": "circular_depth_median", "size": "circular_items"}).assign(week=lambda d: pd.to_datetime(d.week)), on="week", how="left")
    q["access_affordability"] = np.nan
    # illustrative composites (z-scores over the 11 weeks)
    z = lambda s: (s - s.mean()) / s.std(ddof=0) if s.std(ddof=0) > 0 else s * 0
    q["composite_equal_illustrative"] = (z(q.comparable_price_index) + z(q.promo_share) * -1 + z(q.dispersion_median_gap) + z(q.basket_K1)) / 4
    q["composite_alt_illustrative"] = 0.5 * z(q.comparable_price_index) + 0.5 * z(q.basket_K1)
    q.to_csv(P / "qgpai_components.csv", index=False)
    out["components"] = q.round(4).to_dict(orient="records"); out["note"] = "component 5 not computed; composites are illustrations with prespecified weights, not estimates"
    (E / "qgpai_summary.json").write_text(json.dumps(out, indent=2, default=str))
    log.info("\n%s", q[["week", "comparable_price_index", "dispersion_median_gap", "promo_share", "basket_K1", "basket_K2", "basket_unconstrained", "circular_depth_median"]].round(3).to_string(index=False))
    ledger("5", "qgpai", "components written")

if __name__ == "__main__":
    main()