"""Model D — price indices (H5). Preregistration §2 H5. No expenditure weights exist: all indices are unweighted and are
NOT consumer price indices. Size-label history is not observed (only current state) → within-listing package changes are
unobservable; this is reported as a limitation, not estimated.
"""
import json, numpy as np, pandas as pd
from src.common import path, ROOT, get_logger, ledger, seed_everything, SEED
log = get_logger("modelD"); P = path("processed"); E = ROOT / "experiments"

def jevons_chained(df, col):
    """Chained Jevons (matched-model): geometric mean of week-on-week ratios over listings present in both weeks."""
    weeks = sorted(df.week.unique()); idx = [100.0]; n = [np.nan]
    piv = df.pivot_table(index="uid", columns="week", values=col)
    for a, b in zip(weeks, weeks[1:]):
        both = piv[[a, b]].dropna(); r = np.exp((np.log(both[b]) - np.log(both[a])).mean()) if len(both) else np.nan
        idx.append(idx[-1] * r); n.append(len(both))
    return pd.DataFrame({"week": weeks, "index": idx, "n_matched": n})

def fixed_basket(df, col):
    """Balanced listings (present in every week), equal weights: Dutot (ratio of means), Carli (mean of ratios), Jevons (geo mean of ratios)."""
    weeks = sorted(df.week.unique()); piv = df.pivot_table(index="uid", columns="week", values=col).dropna()
    base = piv[weeks[0]]
    rows = []
    for w in weeks:
        rows.append(dict(week=w, n=len(piv), dutot=100 * piv[w].mean() / base.mean(), carli=100 * (piv[w] / base).mean(), jevons=100 * np.exp(np.log(piv[w] / base).mean())))
    return pd.DataFrame(rows)

def naive_mean(df, col):
    g = df.groupby("week")[col].agg(["mean", "median", "size"]).reset_index(); g["index_mean"] = 100 * g["mean"] / g["mean"].iloc[0]; g["index_median"] = 100 * g["median"] / g["median"].iloc[0]
    return g

def main():
    seed_everything(); out = {}
    pw = pd.read_parquet(P / "panel_weekly.parquet"); sf = pw[pw.channel == "online_storefront"].copy()
    res = {}
    for name, d, col in [("storefront_posted", sf, "p_mode"), ("storefront_unit", sf[sf.unit_p.notna()], "unit_p"), ("all_channels_posted", pw, "p_mode")]:
        nm = naive_mean(d, col); jc = jevons_chained(d, col); fb = fixed_basket(d, col)
        m = nm.merge(jc, on="week").merge(fb, on="week", suffixes=("", "_fb")); m["sample"] = name; res[name] = m
        out[name] = dict(naive_mean_change=float(nm.index_mean.iloc[-1] / 100 - 1), naive_median_change=float(nm.index_median.iloc[-1] / 100 - 1), jevons_chained_change=float(jc["index"].iloc[-1] / 100 - 1),
                         fixed_jevons_change=float(fb.jevons.iloc[-1] / 100 - 1), fixed_dutot_change=float(fb.dutot.iloc[-1] / 100 - 1), fixed_carli_change=float(fb.carli.iloc[-1] / 100 - 1), n_balanced=int(fb.n.iloc[0]), n_total_first_week=int(nm["size"].iloc[0]))
    pd.concat(res.values()).to_csv(P / "modelD_indices.csv", index=False)
    # by category (storefront posted)
    bc = []
    for c, d in sf.groupby("category"):
        if d.uid.nunique() < 100: continue
        nm = naive_mean(d, "p_mode"); jc = jevons_chained(d, "p_mode"); fb = fixed_basket(d, "p_mode")
        bc.append(dict(category=c, listings=d.uid.nunique(), naive_mean_change=nm.index_mean.iloc[-1] / 100 - 1, jevons_chained_change=jc["index"].iloc[-1] / 100 - 1, fixed_jevons_change=fb.jevons.iloc[-1] / 100 - 1, n_balanced=fb.n.iloc[0]))
    pd.DataFrame(bc).to_csv(P / "modelD_indices_by_category.csv", index=False)
    # by banner (storefront posted): Jevons chained
    bb = []
    for b, d in sf.groupby("banner"):
        if d.uid.nunique() < 100: continue
        jc = jevons_chained(d, "p_mode"); fb = fixed_basket(d, "p_mode"); nm = naive_mean(d, "p_mode")
        bb.append(dict(banner=b, listings=d.uid.nunique(), jevons_chained_change=jc["index"].iloc[-1] / 100 - 1, fixed_jevons_change=fb.jevons.iloc[-1] / 100 - 1, naive_mean_change=nm.index_mean.iloc[-1] / 100 - 1, n_balanced=fb.n.iloc[0]))
    pd.DataFrame(bb).to_csv(P / "modelD_indices_by_banner.csv", index=False)
    # Bootstrap (listing-level, B=999) of the naive-vs-Jevons gap, storefront posted
    rng = np.random.default_rng(SEED); uids = sf.uid.unique(); piv = sf.pivot_table(index="uid", columns="week", values="p_mode"); weeks = list(piv.columns)
    gaps = []; jev = []
    for b in range(999):
        s = rng.choice(uids, size=len(uids), replace=True); pb = piv.loc[s]
        naive = pb[weeks[-1]].mean() / pb[weeks[0]].mean() - 1
        lj = 0.0
        for a, c in zip(weeks, weeks[1:]):
            both = pb[[a, c]].dropna(); lj += (np.log(both[c]) - np.log(both[a])).mean()
        jev.append(np.exp(lj) - 1); gaps.append(naive - (np.exp(lj) - 1))
    out["bootstrap"] = dict(gap_mean=float(np.mean(gaps)), gap_ci95=[float(np.percentile(gaps, 2.5)), float(np.percentile(gaps, 97.5))], jevons_ci95=[float(np.percentile(jev, 2.5)), float(np.percentile(jev, 97.5))])
    # Assortment turnover (entry/exit) — the observable part of H5
    first = sf.groupby("uid").week.min(); last = sf.groupby("uid").week.max(); weeks_s = sorted(sf.week.unique())
    out["assortment"] = dict(listings=int(sf.uid.nunique()), balanced=int(((first == weeks_s[0]) & (last == weeks_s[-1])).sum()), entrants=int((first > weeks_s[0]).sum()), exits=int((last < weeks_s[-1]).sum()),
                             weekly_entry_rate=float(sf.groupby("week").uid.nunique().pipe(lambda s: (first.value_counts().reindex(weeks_s).fillna(0)[1:] / s.values[:-1]).mean())))
    # Unit-price vs posted: cross-banner unit price gap for matched groups is covered in Model A; here, size-missing share
    out["size_label_history_available"] = False
    # StatCan benchmark: Quebec CPI 'Food purchased from stores' Aug->Sep 2026 ; our storefront Jevons Aug (weeks 13,20,27 Aug) -> Sep (3,10,17,24 Sep)
    cpi = pd.read_csv(P / "statcan_cpi_food.csv"); q = cpi[(cpi.geo == "Quebec") & (cpi["product"] == "Food purchased from stores")].sort_values("ref_date")
    aug = q[q.ref_date.str.startswith("2026-08")].value.iloc[0]; sep = q[q.ref_date.str.startswith("2026-09")].value.iloc[0]
    jc = jevons_chained(sf, "p_mode"); jc["month"] = pd.to_datetime(jc.week).dt.month
    ours = jc.groupby("month")["index"].mean(); out["statcan"] = dict(cpi_qc_food_stores_aug=float(aug), sep=float(sep), cpi_change_aug_sep=float(sep / aug - 1), ours_jevons_aug_sep=float(ours.loc[9] / ours.loc[8] - 1),
                                                                     cpi_latest=str(q.ref_date.iloc[-1]), cpi_yoy_sep=float(sep / q[q.ref_date.str.startswith("2025-09")].value.iloc[0] - 1))
    (E / "modelD_summary.json").write_text(json.dumps(out, indent=2, default=float))
    s = out["storefront_posted"]
    log.info("D storefront: naive mean %.3f, naive median %.3f, Jevons chained %.3f (CI %s), fixed Jevons %.3f (n=%d), Dutot %.3f, Carli %.3f; gap naive−Jevons %.3f CI %s",
             s["naive_mean_change"], s["naive_median_change"], s["jevons_chained_change"], np.round(out["bootstrap"]["jevons_ci95"], 3), s["fixed_jevons_change"], s["n_balanced"], s["fixed_dutot_change"], s["fixed_carli_change"], out["bootstrap"]["gap_mean"], np.round(out["bootstrap"]["gap_ci95"], 3))
    log.info("StatCan QC food-from-stores Aug→Sep %.3f vs ours %.3f ; CPI yoy Sep %.3f", out["statcan"]["cpi_change_aug_sep"], out["statcan"]["ours_jevons_aug_sep"], out["statcan"]["cpi_yoy_sep"])
    ledger("4", "model_d", f"Jevons {s['jevons_chained_change']:.3f} vs naive {s['naive_mean_change']:.3f}")

if __name__ == "__main__":
    main()