"""Model B — promotion dynamics (H2, H3). Preregistration §2 H2–H3.

Storefront promotion state is inferred: week promoted if p_mode <= 0.90 x reference, reference = modal p_mode over the
trailing 8 weeks (incl. current). Circular items carry observed regular prices (original_price) and are analysed separately.
"""
import json, numpy as np, pandas as pd, pyfixest as pf
from scipy import stats
from src.common import path, ROOT, get_logger, ledger, seed_everything, SEED
log = get_logger("modelB"); P = path("processed"); E = ROOT / "experiments"

def rolling_mode(vals, window=8):
    out = []
    for i in range(len(vals)):
        w = vals[max(0, i - window + 1): i + 1]
        vc = pd.Series(w).value_counts()
        top = vc[vc == vc.max()].index
        out.append(max(top))  # ties -> highest price (reference = regular)
    return out

def main():
    seed_everything(); out = {}
    pw = pd.read_parquet(P / "panel_weekly.parquet")
    sf = pw[pw.channel == "online_storefront"].sort_values(["uid", "week"]).copy()
    nw = sf.groupby("uid").week.transform("nunique"); sf = sf[nw >= 8].copy()
    log.info("H2/H3 storefront sample: %d listing-weeks, %d listings, %d banners", len(sf), sf.uid.nunique(), sf.banner.nunique())
    sf["ref"] = sf.groupby("uid").p_mode.transform(lambda s: pd.Series(rolling_mode(s.tolist()), index=s.index))
    sf["promo"] = sf.p_mode <= 0.90 * sf.ref
    sf["depth"] = np.where(sf.promo, np.log(sf.ref / sf.p_mode), 0.0)
    sf["log_ref"] = np.log(sf.ref)
    sf["dlog_p"] = sf.groupby("uid").log_p.diff(); sf["dlog_ref"] = sf.groupby("uid").log_ref.diff()
    sf["promo_lag"] = sf.groupby("uid").promo.shift(1)
    # validation vs current on_sale flag (listings active now; last complete week)
    last = sf[sf.week == sf.week.max()]; last = last[last.active == 1]
    agree = (last.promo == (last.on_sale_now == 1)).mean()
    out["promo_vs_on_sale_now_agreement"] = float(agree); out["promo_rate_last_week"] = float(last.promo.mean()); out["on_sale_now_rate"] = float((last.on_sale_now == 1).mean())
    out["promo_share_listing_weeks"] = float(sf.promo.mean()); out["promo_depth_median"] = float(sf.loc[sf.promo, "depth"].median()); out["promo_depth_mean"] = float(sf.loc[sf.promo, "depth"].mean())
    # H2: volatility raw vs reference
    chg = sf.dropna(subset=["dlog_p"])
    vol = chg.groupby("banner").agg(n=("dlog_p", "size"), sd_raw=("dlog_p", "std"), sd_ref=("dlog_ref", "std"), freq_change=("dlog_p", lambda s: (s.abs() > 1e-9).mean()),
                                    freq_ref_change=("dlog_ref", lambda s: (s.abs() > 1e-9).mean()),
                                    share_changes_promo=("promo", lambda s: np.nan)).reset_index()
    chg["is_change"] = chg.dlog_p.abs() > 1e-9; chg["promo_transition"] = chg.promo != chg.promo_lag
    vol["share_changes_promo"] = vol.banner.map(chg[chg.is_change].groupby("banner").promo_transition.mean())
    vol["ratio_raw_ref"] = vol.sd_raw / vol.sd_ref; vol.sort_values("n", ascending=False).to_csv(P / "modelB_volatility_by_banner.csv", index=False)
    volc = chg.groupby("category").agg(n=("dlog_p", "size"), sd_raw=("dlog_p", "std"), sd_ref=("dlog_ref", "std")).reset_index(); volc["ratio"] = volc.sd_raw / volc.sd_ref
    volc.to_csv(P / "modelB_volatility_by_category.csv", index=False)
    out["sd_raw_all"] = float(chg.dlog_p.std()); out["sd_ref_all"] = float(chg.dlog_ref.std()); out["ratio_all"] = out["sd_raw_all"] / out["sd_ref_all"]
    out["share_weekly_changes_promo_transition"] = float(chg[chg.is_change].promo_transition.mean()); out["freq_weekly_change"] = float(chg.is_change.mean())
    # bootstrap over listings for the ratio (moment aggregation: resample listings with multinomial weights)
    rng = np.random.default_rng(SEED)
    mom = chg.groupby("uid").agg(n=("dlog_p", "size"), s1=("dlog_p", "sum"), q1=("dlog_p", lambda x: (x ** 2).sum()), s2=("dlog_ref", "sum"), q2=("dlog_ref", lambda x: (x ** 2).sum()))
    M = mom.values; rat = []
    for b in range(999):
        w = rng.multinomial(len(M), np.full(len(M), 1 / len(M)))
        n = (w * M[:, 0]).sum(); s1 = (w * M[:, 1]).sum(); q1 = (w * M[:, 2]).sum(); s2 = (w * M[:, 3]).sum(); q2 = (w * M[:, 4]).sum()
        sd1 = np.sqrt(q1 / n - (s1 / n) ** 2); sd2 = np.sqrt(q2 / n - (s2 / n) ** 2); rat.append(sd1 / sd2)
    out["ratio_ci95"] = [float(np.percentile(rat, 2.5)), float(np.percentile(rat, 97.5))]
    # Falsification: single-store online shops (Shopify/WooCommerce) vs anti-bot storefronts — share of multi-price weeks
    single = ["giant_tiger", "pa", "epipresto", "nuvo", "boite_a_grains", "bocoboco", "aliments_merci", "akhavan", "mayrand", "aubut", "maturin", "tau", "loco", "avril"]
    out["multi_price_week_share_single_store"] = float(sf[sf.source.isin(single)].multi_price_week.mean()); out["multi_price_week_share_antibot"] = float(sf[~sf.source.isin(single)].multi_price_week.mean())
    # Transition matrix
    tm = pd.crosstab(sf.promo_lag.dropna().astype(bool), sf.loc[sf.promo_lag.notna(), "promo"], normalize="index"); tm.to_csv(P / "modelB_transition_matrix.csv")
    out["P_stay_promo"] = float(tm.loc[True, True]) if True in tm.index else None; out["P_enter_promo"] = float(tm.loc[False, True]) if False in tm.index else None
    # Episodes
    eps = []
    for uid, g in sf.groupby("uid"):
        g = g.reset_index(drop=True); pr = g.promo.values; i = 0
        while i < len(pr):
            if pr[i]:
                j = i
                while j + 1 < len(pr) and pr[j + 1]: j += 1
                censored_start = i == 0; censored_end = j == len(pr) - 1
                pre_ref = g.ref.iloc[i - 1] if i > 0 else np.nan; post_ref = g.ref.iloc[j + 1] if j + 1 < len(pr) else np.nan
                post_p = g.p_mode.iloc[j + 1] if j + 1 < len(pr) else np.nan
                eps.append(dict(uid=uid, banner=g.banner.iloc[0], category=g.category.iloc[0], start=g.week.iloc[i], end=g.week.iloc[j], duration=j - i + 1,
                                depth=g.depth.iloc[i:j + 1].mean(), censored=censored_start or censored_end, pre_ref=pre_ref, post_ref=post_ref, post_p=post_p,
                                dlog_ref=np.log(post_ref / pre_ref) if pre_ref and post_ref and not (censored_start or censored_end) else np.nan,
                                dlog_post_price_vs_pre_ref=np.log(post_p / pre_ref) if pre_ref and post_p and not (censored_start or censored_end) else np.nan))
                i = j + 1
            else: i += 1
    ep = pd.DataFrame(eps); ep.to_parquet(P / "modelB_episodes.parquet", index=False)
    unc = ep[~ep.censored]
    out["n_episodes"] = int(len(ep)); out["n_uncensored"] = int(len(unc)); out["duration_median_unc"] = float(unc.duration.median()); out["duration_mean_unc"] = float(unc.duration.mean())
    out["share_duration_1w"] = float((unc.duration == 1).mean()); out["depth_median_unc"] = float(unc.depth.median())
    rec = ep.sort_values(["uid", "start"]); rec["prev_end"] = rec.groupby("uid").end.shift(1); rec["recur_weeks"] = (rec.start - rec.prev_end).dt.days / 7
    out["recurrence_median_weeks"] = float(rec.recur_weeks.dropna().median()) if rec.recur_weeks.notna().any() else None; out["share_listings_with_2plus_episodes"] = float((ep.groupby("uid").size() >= 2).mean())
    # Episode classification (A-003): 'return' if post-episode price within ±2 % of the pre-episode reference,
    # 'reset_down' if below, 'reset_up' if above — temporary sales vs regular-price changes that the 10 % rule also flags
    unc = unc.copy(); x = unc.dlog_post_price_vs_pre_ref
    unc["kind"] = np.where(x.abs() <= 0.02, "return", np.where(x < -0.02, "reset_down", "reset_up")); ep.loc[unc.index, "kind"] = unc.kind
    ep.to_parquet(P / "modelB_episodes.parquet", index=False)
    out["episode_kind_shares"] = unc.kind.value_counts(normalize=True).round(4).to_dict(); out["episode_kind_counts"] = unc.kind.value_counts().to_dict()
    ret_ = unc[unc.kind == "return"]
    out["return_episodes"] = dict(n=int(len(ret_)), duration_median=float(ret_.duration.median()), duration_mean=float(ret_.duration.mean()), depth_median=float(ret_.depth.median()),
                                  share_duration_1w=float((ret_.duration == 1).mean()))
    # H3 (temporary sales only): does the reference AFTER the sale differ from the reference BEFORE?  (return class, by construction |Δ price| <= 2 %)
    dr = ret_.dlog_ref.dropna(); out["H3_return_dlog_ref_mean"] = float(dr.mean()); out["H3_return_dlog_ref_se"] = float(dr.std() / np.sqrt(len(dr))); out["H3_return_n"] = int(len(dr))
    out["H3_return_share_up"] = float((dr > 1e-9).mean()); out["H3_return_share_down"] = float((dr < -1e-9).mean())
    # H3 (all low-price episodes): reference change — dominated by regular-price resets
    d = unc.dlog_ref.dropna(); out["H3_dlog_ref_mean"] = float(d.mean()); out["H3_dlog_ref_se"] = float(d.std() / np.sqrt(len(d))); out["H3_n"] = int(len(d))
    out["H3_share_ref_up"] = float((d > 1e-9).mean()); out["H3_share_ref_down"] = float((d < -1e-9).mean()); out["H3_share_ref_same"] = float((d.abs() <= 1e-9).mean())
    d2 = unc.dlog_post_price_vs_pre_ref.dropna(); out["H3_post_price_vs_pre_ref_mean"] = float(d2.mean()); out["H3_post_price_vs_pre_ref_share_up"] = float((d2 > 1e-9).mean())
    # Event study around episode end (weeks -4..+4), listing FE + week FE; pretrend joint test — vectorised merge
    ks = pd.DataFrame({"k": range(-4, 5)})
    evb = ret_[["uid", "start"]].merge(ks, how="cross"); evb["week"] = evb.start + pd.to_timedelta(evb.k * 7, unit="D")   # anchored at the first promoted week of 'return' episodes
    ev = evb.merge(sf[["uid", "week", "log_p", "log_ref"]], on=["uid", "week"], how="inner")[["uid", "week", "k", "log_p", "log_ref"]]
    if len(ev):
        ev["k_f"] = ev.k.astype(int)
        fit = pf.feols("log_p ~ i(k_f, ref=-1) | uid + week", data=ev, vcov={"CRV1": "uid"})
        es = pd.DataFrame({"term": fit.coef().index, "coef": fit.coef().values, "se": fit.se().values}); es.to_csv(P / "modelB_event_study.csv", index=False)
        pre = [t for t in fit.coef().index if any(t.endswith(f"::{k}") or t.endswith(f"[T.{k}]") for k in (-4, -3, -2))]
        try:
            w = fit.wald_test(R=None, q=None) if False else None
        except Exception: w = None
        b = fit.coef()[pre].values; V = fit._vcov; idx = [list(fit.coef().index).index(t) for t in pre]; Vs = V[np.ix_(idx, idx)]
        W = float(b @ np.linalg.solve(Vs, b)); pval = float(1 - stats.chi2.cdf(W, len(pre)))
        out["pretrend_wald"] = W; out["pretrend_p"] = pval; out["event_study_n"] = int(fit._N)
        fitr = pf.feols("log_ref ~ i(k_f, ref=-1) | uid + week", data=ev, vcov={"CRV1": "uid"})
        # Falsification of the pre-sale elevation: single-store online shops (no store-mixture noise) vs anti-bot storefronts
        single = ["giant_tiger", "pa", "epipresto", "nuvo", "boite_a_grains", "bocoboco", "aliments_merci", "akhavan", "mayrand", "aubut", "maturin", "tau", "loco", "avril"]
        ev2 = ev.merge(sf[["uid", "source"]].drop_duplicates(), on="uid"); ev2["feedtype"] = np.where(ev2.source.isin(single), "single_store", "antibot")
        rows = []
        for ft, d in ev2.groupby("feedtype"):
            f2 = pf.feols("log_p ~ i(k_f, ref=-1) | uid + week", data=d, vcov={"CRV1": "uid"})
            for t in f2.coef().index: rows.append(dict(feedtype=ft, term=t, coef=float(f2.coef()[t]), se=float(f2.se()[t]), n=int(f2._N)))
        pd.DataFrame(rows).to_csv(P / "modelB_event_study_by_feedtype.csv", index=False)
        out["event_study_by_feedtype_n"] = ev2.groupby("feedtype").size().to_dict()
        pd.DataFrame({"term": fitr.coef().index, "coef": fitr.coef().values, "se": fitr.se().values}).to_csv(P / "modelB_event_study_ref.csv", index=False)
    # Placebo: random pseudo-episode ends for non-promoted listings
    rng = np.random.default_rng(SEED); pl = []
    nonp = sf[~sf.promo].groupby("uid").filter(lambda g: len(g) >= 9)
    for uid, g in list(nonp.groupby("uid"))[:3000]:
        g = g.sort_values("week"); i = rng.integers(2, len(g) - 2); pre = g.ref.iloc[i - 1]; post = g.ref.iloc[i + 1]
        pl.append(np.log(post / pre))
    pl = np.array(pl); out["placebo_dlog_ref_mean"] = float(pl.mean()); out["placebo_dlog_ref_se"] = float(pl.std() / np.sqrt(len(pl))); out["placebo_n"] = int(len(pl))
    # Circular channel: observed regular prices
    circ = pd.read_parquet(P / "products.parquet"); ret = pd.read_parquet(P / "retailers.parquet")
    circ = circ.merge(ret[["feed_id", "channel", "banner", "is_grocery"]], left_on="source", right_on="feed_id")
    circ = circ[circ.channel.str.startswith("circular") & circ.is_grocery & circ.price.notna()]
    circ["depth_obs"] = np.log(circ.regular_price / circ.price)
    out["circular_items"] = int(len(circ)); out["circular_share_with_regular"] = float(circ.regular_price.notna().mean())
    out["circular_depth_median"] = float(circ.depth_obs.dropna().median()); out["circular_depth_mean"] = float(circ.depth_obs.dropna().mean())
    out["circular_share_recurring_items"] = float((circ.span_days > 8).mean())
    circ.groupby("banner").agg(items=("uid", "size"), with_regular=("regular_price", lambda s: s.notna().mean()), depth_median=("depth_obs", "median"), recurring=("span_days", lambda s: (s > 8).mean())).reset_index().to_csv(P / "modelB_circulars_by_banner.csv", index=False)
    (E / "modelB_summary.json").write_text(json.dumps(out, indent=2, default=float))
    log.info("H2: sd raw %.4f vs ref %.4f (ratio %.2f, CI %s); promo share %.3f depth med %.3f; agreement w/ on_sale %.2f", out["sd_raw_all"], out["sd_ref_all"], out["ratio_all"], np.round(out["ratio_ci95"], 2), out["promo_share_listing_weeks"], out["promo_depth_median"], agree)
    log.info("H3: episodes %d (uncensored %d) kinds %s; return episodes: n %d, duration med %.1f, depth med %.3f, dlog_ref mean %.4f (se %.4f) up %.3f down %.3f; all: dlog_ref mean %.4f; placebo %.4f; pretrend p %.3f (n=%d)",
             out["n_episodes"], out["n_uncensored"], out["episode_kind_shares"], out["return_episodes"]["n"], out["return_episodes"]["duration_median"], out["return_episodes"]["depth_median"],
             out["H3_return_dlog_ref_mean"], out["H3_return_dlog_ref_se"], out["H3_return_share_up"], out["H3_return_share_down"], out["H3_dlog_ref_mean"], out["placebo_dlog_ref_mean"], out.get("pretrend_p", float("nan")), out.get("event_study_n", 0))
    ledger("4", "model_b", f"ratio {out['ratio_all']:.2f}; episodes {out['n_episodes']}; H3 mean {out['H3_dlog_ref_mean']:.4f}")

if __name__ == "__main__":
    main()