"""H5 — Cross-merchant price dispersion of matched products (cross-section at the snapshot).

Sample S_M: match groups (tiers A+B primary; A+B+C sensitivity) with comparable variant structure, both stores in S_R
(region known), stores not siblings, groups spanning >= 2 stores after filters.
Model: log P_im = alpha_i + gamma_r(m) + delta_platform(m) + b1 log n_m + b2 Reseller_m + b3 Certified_m + e_im,
group fixed effects absorbed by within-group demeaning; SE clustered by store; joint test of region effects and the
t-test on Reseller use a wild cluster bootstrap (Rademacher, 9 999 draws, restricted residuals).
Outputs experiments/h5_hedonic.json, data/processed/h5_groups.parquet, data/processed/h5_region_effects.csv
"""
from __future__ import annotations
import json
import numpy as np, pandas as pd
from src.common import ROOT, P, SEED, get_logger, ledger, register_output, write_json, PROC_DIR, EXP_DIR

log = get_logger("models.h5")
PROC = PROC_DIR
B = 9999


def build_sample(tiers):
    pr = pd.read_parquet(P["interim"] / "match_pairs.parquet")
    li = pd.read_parquet(P["interim"] / "listings_clean.parquet").set_index("uid")
    sp = pd.read_parquet(PROC / "store_panel.parquet").set_index("id")
    sib = pd.read_parquet(P["interim"] / "sibling_stores.parquet").set_index("id").sibling_rep
    pr = pr[pr.tier.isin(tiers) & pr.comparable]
    ok = lambda s: s.map(sp.region_final).fillna("Inconnue") != "Inconnue"
    pr = pr[ok(pr.store_a) & ok(pr.store_b)]
    pr = pr[pr.store_a.map(sib).fillna(pr.store_a) != pr.store_b.map(sib).fillna(pr.store_b)]
    import networkx as nx
    G = nx.Graph(); G.add_edges_from(pr[["uid_a", "uid_b"]].itertuples(index=False, name=None))
    rows = [(g, u) for g, comp in enumerate(nx.connected_components(G)) for u in comp]
    d = pd.DataFrame(rows, columns=["group", "uid"])
    d["store"] = d.uid.map(li.store_id); d["price"] = d.uid.map(li.price); d["category"] = d.uid.map(li.category)
    d["region"] = d.store.map(sp.region_final); d["platform"] = d.store.map(sp.platform).fillna("other")
    d["log_n"] = np.log(d.store.map(sp.n_listings)); d["reseller"] = d.store.map(sp.store_kind).isin(["revendeur", "collectif"]).astype(float)
    d["certified"] = d.store.map(sp.certified).fillna(False).astype(float); d["lp"] = np.log(d.price)
    gs = d.groupby("group").store.nunique(); d = d[d.group.map(gs) >= 2]
    return d


def within(df, cols):
    return df[cols] - df.groupby("group")[cols].transform("mean")


def fit(d, rng, label):
    # design: region dummies (drop the largest region), platform dummies (drop shopify), controls
    regions = sorted(d.region.unique()); base_r = d.region.value_counts().idxmax()
    plats = sorted(d.platform.unique()); base_p = "shopify" if "shopify" in plats else plats[0]
    X = pd.DataFrame(index=d.index)
    for r in regions:
        if r != base_r: X[f"r::{r}"] = (d.region == r).astype(float)
    for p_ in plats:
        if p_ != base_p: X[f"p::{p_}"] = (d.platform == p_).astype(float)
    X["log_n"] = d.log_n; X["reseller"] = d.reseller; X["certified"] = d.certified
    cols = list(X.columns); dd = pd.concat([d[["group", "store", "lp"]], X], axis=1)
    W = within(dd, cols + ["lp"]); Xw = W[cols].values; yw = W["lp"].values
    keep = Xw.std(0) > 1e-10; Xw = Xw[:, keep]; cols = [c for c, k in zip(cols, keep) if k]
    beta, *_ = np.linalg.lstsq(Xw, yw, rcond=None); e = yw - Xw @ beta
    cl = d.store.astype("category").cat.codes.values; G = cl.max() + 1
    XtX_inv = np.linalg.pinv(Xw.T @ Xw)
    def cluster_cov(Xw, e):
        meat = np.zeros((Xw.shape[1], Xw.shape[1]))
        for g in range(G):
            m = cl == g; s = Xw[m].T @ e[m]; meat += np.outer(s, s)
        return XtX_inv @ meat @ XtX_inv * G / (G - 1)
    V = cluster_cov(Xw, e); se = np.sqrt(np.diag(V))
    ridx = [i for i, c in enumerate(cols) if c.startswith("r::")]
    def wald(beta, V, idx):
        b = beta[idx]; Vi = V[np.ix_(idx, idx)]; return float(b @ np.linalg.pinv(Vi) @ b)
    W_obs = wald(beta, V, ridx)
    # wild cluster bootstrap: restricted model (no region dummies) residuals, Rademacher by store
    nr = [i for i in range(len(cols)) if i not in ridx]
    b_r, *_ = np.linalg.lstsq(Xw[:, nr], yw, rcond=None); e_r = yw - Xw[:, nr] @ b_r; yhat_r = Xw[:, nr] @ b_r
    Wb = np.empty(B)
    i_res = cols.index("reseller") if "reseller" in cols else None
    nr2 = [i for i in range(len(cols)) if i != i_res] if i_res is not None else None
    if i_res is not None:
        b_r2, *_ = np.linalg.lstsq(Xw[:, nr2], yw, rcond=None); e_r2 = yw - Xw[:, nr2] @ b_r2; yhat_r2 = Xw[:, nr2] @ b_r2
    tb = np.empty(B); t_obs = beta[i_res] / se[i_res] if i_res is not None else np.nan
    gcode = d.group.astype("category").cat.codes.values; ng = gcode.max() + 1; gsize = np.bincount(gcode)
    def gdemean(v):   # the match-group fixed effects are absorbed by within-transformation: the bootstrap disturbance
        return v - (np.bincount(gcode, weights=v) / gsize)[gcode]   # must be re-projected off the group means (referee 1, must-fix)
    for b in range(B):
        w = rng.choice([-1.0, 1.0], G)[cl]
        yb = yhat_r + gdemean(e_r * w); bb = XtX_inv @ (Xw.T @ yb); eb = yb - Xw @ bb
        Wb[b] = wald(bb, cluster_cov(Xw, eb), ridx)
        if i_res is not None:
            yb2 = yhat_r2 + gdemean(e_r2 * w); bb2 = XtX_inv @ (Xw.T @ yb2); eb2 = yb2 - Xw @ bb2; V2 = cluster_cov(Xw, eb2)
            tb[b] = bb2[i_res] / np.sqrt(V2[i_res, i_res])
    p_region = float((np.sum(Wb >= W_obs) + 1) / (B + 1))
    p_res = float((np.sum(np.abs(tb) >= abs(t_obs)) + 1) / (B + 1)) if i_res is not None else None
    # leave-one-region-out joint Wald (which region drives the rejection?) and single-region Wald contributions
    loo = {}
    for r in regions:
        if r == base_r: continue
        keep_idx = [i for i in ridx if cols[i] != f"r::{r}"]
        loo[r] = round(wald(beta, V, keep_idx), 2)
    contrib = {cols[i][3:]: round(float(beta[i] ** 2 / V[i, i]), 2) for i in ridx}
    n_groups = int(d.group.nunique())
    gaps = d.groupby("group").lp.agg(lambda s: s.max() - s.min())
    reg_stores = d.groupby("region").store.nunique()
    res = {"label": label, "obs": int(len(d)), "groups": n_groups, "stores": int(G), "regions": len(regions), "base_region": base_r,
           "regions_with_ge5_stores": int((reg_stores >= 5).sum()), "stores_by_region": reg_stores.to_dict(),
           "coef": {c: {"b": round(float(b), 4), "se_cluster": round(float(s), 4)} for c, b, s in zip(cols, beta, se)},
           "wald_region": round(W_obs, 3), "df_region": len(ridx), "p_region_wildboot": p_region, "t_reseller": round(float(t_obs), 3) if i_res is not None else None, "p_reseller_wildboot": p_res,
           "within_R2": round(float(1 - e.var() / yw.var()), 4), "wald_loo_region": loo, "wald_single_region_contrib": contrib, "groups_spanning_regions": int((d.groupby("group").region.nunique() >= 2).sum()),
           "gap_median": round(float(gaps.median()), 4), "gap_p90": round(float(gaps.quantile(0.9)), 4), "share_gap_gt5pct": round(float((gaps > 0.05).mean()), 4),
           "share_gap_zero": round(float((gaps < 1e-9).mean()), 4), "B": B}
    eff = pd.DataFrame({"region": [c[3:] for c in cols if c.startswith("r::")], "effect": [beta[i] for i in ridx], "se": [se[i] for i in ridx]})
    eff = pd.concat([eff, pd.DataFrame({"region": [base_r], "effect": [0.0], "se": [0.0]})]); eff["label"] = label; eff["n_stores"] = eff.region.map(reg_stores)
    return res, eff, d


def main():
    rng = np.random.default_rng(SEED); out = []; effs = []
    for label, tiers in (("primary_AB", ["A", "B"]), ("sens_ABC", ["A", "B", "C"])):
        d = build_sample(tiers)
        if d.group.nunique() < 20:
            out.append({"label": label, "status": "insufficient", "groups": int(d.group.nunique())}); continue
        res, eff, dd = fit(d, rng, label); out.append(res); effs.append(eff); log.info(json.dumps({k: v for k, v in res.items() if k not in ("coef", "stores_by_region")}))
        if label == "primary_AB":
            dd.to_parquet(PROC / "h5_groups.parquet", index=False); register_output(PROC / "h5_groups.parquet", "h5")
    pd.concat(effs).to_csv(PROC / "h5_region_effects.csv", index=False)
    # descriptive certification contrast (NOT H6): within-category, store-clustered
    li = pd.read_parquet(P["interim"] / "listings_clean.parquet"); sp = pd.read_parquet(PROC / "store_panel.parquet").set_index("id")
    s = li[li.f_analysis_price].copy(); s["cert"] = s.store_id.map(sp.cert_adq | sp.cert_lpq).fillna(False).astype(float)
    s["region"] = s.store_id.map(sp.region_final); s = s[s.region.notna() & (s.region != "Inconnue")]
    s["y"] = s.log_price - s.groupby("category").log_price.transform("mean"); s["x"] = s.cert - s.groupby("category").cert.transform("mean")
    b = float((s.x * s.y).sum() / (s.x ** 2).sum()); r = s.y - b * s.x
    g = s.assign(r=r).groupby("store_id").apply(lambda h: (h.x * h.r).sum()); se = float(np.sqrt((g ** 2).sum()) / (s.x ** 2).sum())
    cert = {"note": "descriptive store-certification contrast (Aliments du Québec / Les Produits du Québec), category FE, store-clustered SE — NOT a test of H6",
            "obs": int(len(s)), "stores": int(s.store_id.nunique()), "certified_stores": int(s.groupby('store_id').cert.first().sum()), "b_log_price": round(b, 4), "se": round(se, 4)}
    write_json({"models": out, "certification_contrast": cert}, EXP_DIR / "h5_hedonic.json", "h5"); ledger("models.h5_hedonic", models=[{k: v for k, v in m.items() if k not in ("coef", "stores_by_region")} for m in out], cert=cert)
    log.info(json.dumps(cert))


if __name__ == "__main__":
    main()