"""Phase 6 — robustness and falsification checks that are not already inside the model scripts.

Runs: balanced common-product panel (H1); unit-price normalisation (H1 on log unit prices); outlier trimming (1/99 %);
category heterogeneity (H1); exclusion of same-company pairs; promotion-week exclusion (H1 gaps on non-promoted weeks only);
Holm correction over the preregistered primary tests; attrition table across stages. Results → data/processed/robustness_*.csv
and experiments/robustness_summary.json.
"""
import json, numpy as np, pandas as pd
from scipy import stats
from src.common import path, ROOT, get_logger, ledger, seed_everything, SEED
log = get_logger("robust"); P = path("processed"); E = ROOT / "experiments"

def gaps(pm, col="log_p"):
    return pm.groupby(["group_id", "week"])[col].agg(lambda s: s.max() - s.min())

def main():
    seed_everything(); 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].copy()
    base = gaps(pm); out["baseline_median_gap"] = float(base.median())
    # 1 balanced panel: groups with >= 2 banners in every one of the 11 weeks
    nwk = pm.groupby("group_id").week.nunique(); bal_groups = nwk[nwk == pm.week.nunique()].index
    bal = pm[pm.group_id.isin(bal_groups)]; gb = gaps(bal)
    out["balanced"] = dict(groups=int(len(bal_groups)), median_gap=float(gb.median()), share_ge_5pct=float((gb >= .05).mean()))
    # 2 unit-price normalisation
    pu = pm[pm.log_unit_p.notna()]; gu = gaps(pu, "log_unit_p"); out["unit_price"] = dict(group_weeks=int(len(gu)), median_gap=float(gu.median()), share_ge_5pct=float((gu >= .05).mean()))
    # 3 trimming
    lo, hi = base.quantile(.01), base.quantile(.99); tr = base[(base >= lo) & (base <= hi)]; out["trimmed_1_99"] = dict(median_gap=float(tr.median()), mean_gap=float(tr.mean()), mean_untrimmed=float(base.mean()))
    # 4 category heterogeneity (with Holm-adjusted test that median gap >= 5 % via sign test on gap-0.05)
    rows = []
    for c, g in pm.groupby("category"):
        gg = gaps(g)
        if len(gg) < 50: continue
        k = int((gg > .05).sum()); n = int(len(gg)); p = stats.binomtest(k, n, 0.5, alternative="greater").pvalue
        rows.append(dict(category=c, group_weeks=n, median_gap=float(gg.median()), share_ge_5pct=float((gg >= .05).mean()), p_sign_test=p))
    cat = pd.DataFrame(rows).sort_values("median_gap", ascending=False)
    m = len(cat); order = np.argsort(cat.p_sign_test.values); adj = np.empty(m)
    for rank, i in enumerate(order): adj[i] = min(1.0, cat.p_sign_test.values[i] * (m - rank))
    cat["p_holm"] = np.maximum.accumulate(adj[order])[np.argsort(order)]; cat.to_csv(P / "robustness_category_gaps.csv", index=False)
    # 5 same-company exclusion (cross-company only) and same-company only
    xc = pm[pm.groupby(["group_id", "week"]).company.transform("nunique") >= 2]; gx = gaps(xc)
    out["cross_company_only"] = dict(group_weeks=int(len(gx)), median_gap=float(gx.median()))
    # 6 promotion exclusion: drop listing-weeks at <= 0.9 x trailing mode
    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)
    pm = pm.sort_values(["uid", "week"]); pm["ref"] = pm.groupby("uid").p_mode.transform(rolling_mode); pm["promo"] = pm.p_mode <= 0.9 * pm.ref
    np_ = pm[~pm.promo]; np_ = np_[np_.groupby(["group_id", "week"]).banner.transform("nunique") >= 2]; gn = gaps(np_)
    out["non_promo_weeks_only"] = dict(group_weeks=int(len(gn)), median_gap=float(gn.median()), share_ge_5pct=float((gn >= .05).mean()))
    po = pm[pm.groupby(["group_id", "week"]).promo.transform("max")]; gp = gaps(po); out["group_weeks_with_a_promo"] = dict(group_weeks=int(len(gp)), median_gap=float(gp.median()))
    # 7 Holm over preregistered primary tests (one p per hypothesis family where a p-value exists)
    A = json.load(open(E / "modelA_summary.json")); prim = {}
    # H1: sign test that median gap > 5 %
    k = int((base > .05).sum()); prim["H1_median_gap_gt_5pct"] = stats.binomtest(k, len(base), 0.5, alternative="greater").pvalue
    prim["H1_AR1_gt_0.5"] = float(1 - stats.norm.cdf((A["ar1_rho"] - 0.5) / A["ar1_se"]))
    if (E / "modelB_summary.json").exists():
        B = json.load(open(E / "modelB_summary.json"))
        lo_, hi_ = B["ratio_ci95"]; prim["H2_ratio_ne_1"] = float(2 * min(np.mean(np.array([lo_, hi_]) <= 1), 0.5)) if not (lo_ <= 1 <= hi_) else 1.0
        prim["H2_ratio_ne_1"] = 0.0 if not (lo_ <= 1 <= hi_) else 1.0  # CI-based (bootstrap): reject at 5 % iff CI excludes 1
        prim["H3_return_episodes_dlog_ref_ne_0"] = float(2 * (1 - stats.norm.cdf(abs(B["H3_return_dlog_ref_mean"] / B["H3_return_dlog_ref_se"]))))
        prim["H3_all_low_price_episodes_dlog_ref_ne_0"] = float(2 * (1 - stats.norm.cdf(abs(B["H3_dlog_ref_mean"] / B["H3_dlog_ref_se"]))))
    if (E / "modelC_summary.json").exists():
        C = json.load(open(E / "modelC_summary.json"))
        c = C["rival_all"]; z = c["coefs"]["rival_all_l1"] / c["ses"]["rival_all_l1"]; prim["H4_lag1_ne_0"] = float(2 * (1 - stats.norm.cdf(abs(z))))
    if (E / "modelD_summary.json").exists():
        D = json.load(open(E / "modelD_summary.json")); lo_, hi_ = D["bootstrap"]["gap_ci95"]; prim["H5_naive_vs_jevons_gap_ne_0"] = 0.0 if not (lo_ <= 0 <= hi_) else 1.0
    if (E / "modelE_summary.json").exists():
        Ee = json.load(open(E / "modelE_summary.json")); prim["H6_K2_captures_all"] = np.nan  # descriptive share; no test
    ps = pd.Series({k: v for k, v in prim.items() if not (isinstance(v, float) and np.isnan(v))}).sort_values()
    holm = pd.DataFrame({"test": ps.index, "p": ps.values}); mm = len(holm); holm["p_holm"] = np.minimum(1, np.maximum.accumulate(holm.p.values * (mm - np.arange(mm))))
    holm.to_csv(P / "robustness_holm.csv", index=False); out["holm"] = holm.to_dict(orient="records")
    (E / "robustness_summary.json").write_text(json.dumps(out, indent=2, default=float))
    log.info("%s", json.dumps({k: v for k, v in out.items() if k != "holm"}, indent=None, default=float)[:900]); log.info("\n%s", holm.to_string(index=False))
    ledger("6", "robustness", "robustness tables written")

if __name__ == "__main__":
    main()