"""H4 — Related diversification (EXPLORATORY). Weekly (store, category) presence reconstructed from first_seen/last_seen.
Baseline weeks W36–W39 (2026-08-31 … 2026-09-27): portfolio and relatedness (proximity across stores) computed here only.
Outcome: entry into a category absent in W36–W39 and present (>= 1 listing) in W40–W44 (2026-09-28 … 2026-11-01).
LPM Entry_sc = b * Density_sc + a_s + g_c + e, SE clustered by store; placebo with category-relabelled proximity.
Region-level entries counted descriptively. Connector-change audit: share of entering stores whose sync_log shows a
status change or an 'added' burst in the outcome window.
"""
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, BASE_PROC, BASE_EXP, LABEL

log = get_logger("models.h4")
PROC = PROC_DIR
BASE = (pd.Timestamp("2026-08-31", tz="UTC"), pd.Timestamp("2026-09-28", tz="UTC"))
OUT = (pd.Timestamp("2026-09-28", tz="UTC"), pd.Timestamp("2026-11-02", tz="UTC"))


def present(li, start, end):
    """listings alive at some point in [start, end): first_seen < end and last_seen >= start"""
    m = (li.first_seen_utc < end) & (li.last_seen_utc >= start)
    return li[m].groupby(["store_id", "category"]).size().rename("n").reset_index()


def proximity(pres):
    pv = pres.assign(v=1).pivot_table(index="store_id", columns="category", values="v", fill_value=0)
    co = pv.T.values @ pv.values; k = np.diag(co).astype(float)
    phi = co / np.maximum(np.minimum.outer(k, k), 1); np.fill_diagonal(phi, 0); return pd.DataFrame(phi, index=pv.columns, columns=pv.columns), pv


def lpm(df, rng):
    # within transformation for two-way FE (store, category) via iterative demeaning
    y = df.entry.values.astype(float); x = df.density.values.astype(float)
    s = df.store_id.astype("category").cat.codes.values; c = df.category.astype("category").cat.codes.values
    def demean(v):
        v = v.copy()
        for _ in range(50):
            v -= pd.Series(v).groupby(s).transform("mean").values; v -= pd.Series(v).groupby(c).transform("mean").values
        return v
    yw, xw = demean(y), demean(x); b = float((xw * yw).sum() / (xw ** 2).sum()); r = yw - b * xw
    g = pd.Series(xw * r).groupby(s).sum(); se = float(np.sqrt((g ** 2).sum()) / (xw ** 2).sum())
    return b, se


def main():
    rng = np.random.default_rng(SEED)
    # all PUBLISHED listings of live stores, including listings inactive at the snapshot (listings_clean keeps active=1 only)
    raw = pd.read_parquet(P["interim"] / "listings_raw.parquet"); raw = raw[raw.listing_status == "published"]
    lc = pd.read_parquet(P["interim"] / "listings_clean.parquet")
    live = set(lc.loc[lc.f_live_store, "store_id"]); bad = set(lc.loc[lc.f_hors_mission | lc.f_sibling_dup, "store_id"])
    li = raw[raw.store_id.isin(live) & ~raw.store_id.isin(bad)].copy()
    if LABEL == "model":
        cm = pd.read_parquet(BASE_PROC / "category_model.parquet").set_index("uid").category_model
        li = li.assign(category=li.uid.map(cm).fillna(li.category))
    sp = pd.read_parquet(PROC / "store_panel.parquet").set_index("id")
    base = present(li, *BASE); out = present(li, *OUT)
    phi, pv = proximity(base)
    cats = list(pv.columns); stores = pv.index
    dens = (pv.values @ phi.values) / np.maximum(phi.values.sum(0), 1e-9)   # density of c around store's baseline portfolio
    D = pd.DataFrame(dens, index=stores, columns=cats).stack().rename("density").reset_index().rename(columns={"level_0": "store_id", "level_1": "category"})
    D = D.merge(pv.stack().rename("base_present").reset_index().rename(columns={"level_0": "store_id", "level_1": "category"}), on=["store_id", "category"])
    outp = out.assign(out_present=1)[["store_id", "category", "out_present"]]
    D = D.merge(outp, on=["store_id", "category"], how="left"); D["out_present"] = D.out_present.fillna(0)
    D = D[D.base_present == 0].copy(); D["entry"] = D.out_present.astype(int)
    b, se = lpm(D, rng)
    # placebo: relabel categories in the proximity matrix (random permutation of rows+columns of phi) so that the
    # store's portfolio is scored against a wrong relatedness structure; 200 draws
    plac = []
    for k in range(200):
        pi = rng.permutation(len(cats)); php = phi.values[np.ix_(pi, pi)]
        dp = (pv.values @ php) / np.maximum(php.sum(0), 1e-9)
        Dp = D.copy(); Dp["density"] = pd.DataFrame(dp, index=stores, columns=cats).stack().reindex(pd.MultiIndex.from_frame(D[["store_id", "category"]])).values
        plac.append(lpm(Dp, rng)[0])
    # region-level entries (descriptive)
    sr = sp.region_final
    rb = base.assign(region=base.store_id.map(sr)).query("region!='Inconnue'").groupby(["region", "category"]).n.sum()
    ro = out.assign(region=out.store_id.map(sr)).query("region!='Inconnue'").groupby(["region", "category"]).n.sum()
    all_rc = pd.MultiIndex.from_product([sorted(sr[sr != "Inconnue"].unique()), cats])
    rb = rb.reindex(all_rc).fillna(0); ro = ro.reindex(all_rc).fillna(0)
    region_entries = [(r, c) for (r, c) in all_rc if rb[(r, c)] == 0 and ro[(r, c)] > 0]
    region_exits = [(r, c) for (r, c) in all_rc if rb[(r, c)] > 0 and ro[(r, c)] == 0]
    # connector audit: entering stores with an 'added' burst (> 50 % of baseline catalogue added in one sync) in the outcome window
    sl = pd.read_parquet(P["interim"] / "sync_log.parquet"); slo = sl[(sl.ts_utc >= OUT[0]) & (sl.ts_utc < OUT[1])]
    basesize = base.groupby("store_id").n.sum()
    burst = slo.groupby("store_id").added.max() / basesize.reindex(slo.store_id.unique()).fillna(1)
    entering = D[D.entry == 1].store_id.unique()
    share_burst = float((burst.reindex(entering).fillna(0) > 0.5).mean()) if len(entering) else None
    inactive_n = int((li.active == 0).sum())
    res = {"status": "exploratory", "listings_used": int(len(li)), "inactive_listings_included": inactive_n, "baseline_window": [str(BASE[0].date()), str((BASE[1] - pd.Timedelta(days=1)).date())], "outcome_window": [str(OUT[0].date()), str((OUT[1] - pd.Timedelta(days=1)).date())],
           "stores": int(pv.shape[0]), "store_category_cells_absent_at_baseline": int(len(D)), "entries": int(D.entry.sum()), "entering_stores": int(len(entering)),
           "lpm_density_coef": round(b, 4), "se_cluster_store": round(se, 4), "t": round(b / se, 3) if se > 0 else None,
           "placebo_coef_mean": round(float(np.mean(plac)), 4), "placebo_p95": round(float(np.percentile(plac, 95)), 4), "placebo_p": float((np.sum(np.array(plac) >= b) + 1) / (len(plac) + 1)),
           "density_mean_entry": round(float(D[D.entry == 1].density.mean()), 4), "density_mean_no_entry": round(float(D[D.entry == 0].density.mean()), 4),
           "region_level_entries": region_entries, "region_level_exits": region_exits, "n_region_entries": len(region_entries),
           "share_entering_stores_with_added_burst": share_burst}
    D.to_parquet(PROC / "h4_panel.parquet", index=False); register_output(PROC / "h4_panel.parquet", "h4")
    write_json(res, EXP_DIR / "h4_dynamics.json", "h4"); ledger("models.h4_dynamics", **{k: v for k, v in res.items() if "region_level" not in k})
    log.info(json.dumps({k: v for k, v in res.items() if "region_level" not in k}))


if __name__ == "__main__":
    main()