"""H1 — Regional specialization: RAA matrix, Cramér's V, store-level region-reshuffle permutation test, store bootstrap CIs.
Preregistered: scheme M primary, decision requires p < 0.05 under M and T (Holm within the confirmatory family, applied in src/models/summary.py).
Outputs experiments/h1_raa.json, data/processed/raa_matrix.csv (point estimates + CIs per sample/scheme).
"""
from __future__ import annotations
import json
import numpy as np, pandas as pd
from src.common import ROOT, P, REGIONS, SEED, get_logger, ledger, register_output, write_json, PROC_DIR, EXP_DIR

log = get_logger("models.h1")
PROC = PROC_DIR
N_PERM, N_BOOT = 2000, 2000
CATS = None


def weights(sc: pd.DataFrame, scheme: str) -> np.ndarray:
    sizes = sc.groupby("store_id").n.transform("sum")
    if scheme == "L": return sc.n.values.astype(float)
    if scheme == "M": return (sc.n / sizes).values
    cut = sc.groupby("store_id").n.sum().quantile(0.99)
    return np.where(sizes > cut, 0.0, sc.n.values.astype(float))


def table(region_codes, cat_codes, w, R, C):
    t = np.zeros((R, C)); np.add.at(t, (region_codes, cat_codes), w); return t


def cramers_v(t):
    n = t.sum(); r = t.sum(1, keepdims=True); c = t.sum(0, keepdims=True)
    e = r @ c / n; m = e > 0
    chi2 = ((t[m] - e[m]) ** 2 / e[m]).sum()
    return float(np.sqrt(chi2 / (n * (min(t.shape) - 1))))


def raa(t):
    n = t.sum(); return (t / t.sum(1, keepdims=True)) / (t.sum(0, keepdims=True) / n)


def run(sc, sample, scheme, rng):
    d = sc[(sc.region_final != "Inconnue")] if sample == "S_R" else sc[sc.region_db.notna()].assign(region_final=lambda x: x.region_dbfirst)
    d = d[d.region_final != "Inconnue"].reset_index(drop=True)
    regions = sorted(d.region_final.unique()); cats = sorted(d.category.unique())
    rcode = d.region_final.map({r: i for i, r in enumerate(regions)}).values; ccode = d.category.map({c: i for i, c in enumerate(cats)}).values
    w = weights(d, scheme); R, C = len(regions), len(cats)
    t = table(rcode, ccode, w, R, C); v = cramers_v(t); A = raa(t)
    # permutation: reshuffle region labels across STORES
    stores = d.store_id.unique(); sreg = d.drop_duplicates("store_id").set_index("store_id").region_final
    sidx = d.store_id.map({s: i for i, s in enumerate(stores)}).values
    srcode = np.array([ {r: i for i, r in enumerate(regions)}[sreg[s]] for s in stores])
    vperm = np.empty(N_PERM)
    for k in range(N_PERM):
        perm = rng.permutation(srcode); vperm[k] = cramers_v(table(perm[sidx], ccode, w, R, C))
    p = float((np.sum(vperm >= v) + 1) / (N_PERM + 1))
    # store bootstrap
    vb = np.empty(N_BOOT); Ab = np.empty((N_BOOT, R, C))
    rows_by_store = [np.where(sidx == i)[0] for i in range(len(stores))]
    for b in range(N_BOOT):
        draw = rng.integers(0, len(stores), len(stores))
        idx = np.concatenate([rows_by_store[i] for i in draw])
        tb = table(rcode[idx], ccode[idx], w[idx], R, C)
        vb[b] = cramers_v(tb); Ab[b] = raa(np.where(tb == 0, 1e-12, tb))
    lo, hi = np.nanpercentile(Ab, 2.5, axis=0), np.nanpercentile(Ab, 97.5, axis=0)
    mat = pd.DataFrame({"sample": sample, "scheme": scheme, "region": np.repeat(regions, C), "category": np.tile(cats, R),
                        "n": t.ravel(), "raa": A.ravel(), "raa_lo": lo.ravel(), "raa_hi": hi.ravel()})
    mat["sig_gt1"] = mat.raa_lo > 1; mat["sig_lt1"] = mat.raa_hi < 1
    res = {"sample": sample, "scheme": scheme, "stores": int(len(stores)), "regions": R, "categories": C, "cramers_v": round(v, 4),
           "v_boot_percentile_note": "store resampling with replacement duplicates stores and biases V upward; not used for inference",
           "v_boot_p2.5_97.5": [round(float(np.percentile(vb, 2.5)), 4), round(float(np.percentile(vb, 97.5)), 4)],
           "perm_null_mean": round(float(vperm.mean()), 4), "perm_null_p5": round(float(np.percentile(vperm, 5)), 4), "excess_v": round(float(v - vperm.mean()), 4),
           "perm_null_p95": round(float(np.percentile(vperm, 95)), 4), "p_perm": p, "cells_sig_gt1": int(mat.sig_gt1.sum()), "cells_sig_lt1": int(mat.sig_lt1.sum()), "cells": R * C}
    return res, mat


def main():
    sc = pd.read_parquet(PROC / "store_category.parquet")
    rng = np.random.default_rng(SEED)
    results, mats = [], []
    for sample in ("S_R", "S_R0"):
        for scheme in ("M", "L", "T"):
            res, mat = run(sc, sample, scheme, rng); results.append(res); mats.append(mat); log.info(json.dumps(res))
    # merchant-count RAA (stores instead of listings) for comparison
    d = sc[sc.region_final != "Inconnue"].drop_duplicates(["store_id", "category"]).assign(n=1)
    res_m, mat_m = run(d, "S_R", "L", rng); res_m["scheme"] = "stores"; mat_m["scheme"] = "stores"
    results.append(res_m); mats.append(mat_m)
    pd.concat(mats).to_csv(PROC / "raa_matrix.csv", index=False); register_output(PROC / "raa_matrix.csv", "h1")
    write_json({"n_perm": N_PERM, "n_boot": N_BOOT, "results": results}, EXP_DIR / "h1_raa.json", "h1"); ledger("models.h1_raa", results=results)


if __name__ == "__main__":
    main()