"""Robustness variants (preregistered §5): H1 Cramér's V with store-reshuffle permutation p, H2-type diversity range and H3 Gini
recomputed on (i) S_R0, (ii) certification/association sources only, (iii) excluding Montréal, (iv) excluding the 1 % largest stores,
(v) 8 macro-categories, (vi) 'autre' dropped, (vii) categories reassigned by the audit confusion matrix (100 draws, if the audit exists),
(viii) excluding resellers/collectives, (ix) region conflicts dropped. Output experiments/robustness.json"""
from __future__ import annotations
import json
import numpy as np, pandas as pd
from src.common import ROOT, SEED, get_logger, ledger, write_json, PROC_DIR, EXP_DIR, P, BASE_PROC, BASE_EXP
from src.models.h1_raa import weights, table, cramers_v
from src.models.h3_longtail import gini

log = get_logger("validation.robustness")
PROC = PROC_DIR
MACRO = {"erable": "food", "epicerie": "food", "cafe_the": "food", "chocolat": "food", "boulangerie": "food", "viande_poisson": "food", "fromage_laitier": "food", "boissons": "food",
         "alcool": "alcohol", "sante_beaute": "beauty", "maison": "home", "mode": "fashion", "bijoux": "fashion", "art": "art", "enfants": "kids_pets", "animaux": "kids_pets",
         "plein_air": "outdoor_garden", "jardin": "outdoor_garden", "industriel": "industrial", "autre": "other"}
CERT_SOURCES = {"alimentsduquebec_ent", "lesproduitsduquebec_ent", "lesproduitsduquebec_det", "cartv_bio", "ctaq", "erabledici", "fromagesdici", "cidreduquebec", "vinsduquebec",
                "distilleriesduquebec", "ambq_profils", "artisansaloeuvre", "afmq", "apiculteursduquebec", "boeufdici", "zoneboreale", "createursdesaveurs", "croquezoutaouais",
                "arretsgourmands", "goutezy", "laurentidesjenmange", "goutcotenord", "gaspesiegourmande", "charlevoix", "mauriciemiam", "gardemangerduquebec", "saveursbsl", "goutezat"}


def h1(d, rng, scheme="M", nperm=1000):
    d = d.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); v = cramers_v(table(rcode, ccode, w, R, C))
    stores = d.store_id.unique(); sidx = d.store_id.map({s: i for i, s in enumerate(stores)}).values
    sreg = d.drop_duplicates("store_id").set_index("store_id").region_final.reindex(stores).map({r: i for i, r in enumerate(regions)}).values
    null = np.array([cramers_v(table(rng.permutation(sreg)[sidx], ccode, w, R, C)) for _ in range(nperm)])
    return {"V": round(v, 4), "null_mean": round(float(null.mean()), 4), "p": float((np.sum(null >= v) + 1) / (nperm + 1)), "stores": int(len(stores)), "regions": R}


def ent_range(d):
    out = {}
    for r, g in d.groupby("region_final"):
        w = g.n / g.groupby("store_id").n.transform("sum"); p = (w.groupby(g.category).sum() / w.sum()).values; p = p[p > 0]
        out[r] = float(-(p * np.log(p)).sum() / np.log(d.category.nunique()))
    s = pd.Series(out); return {"min": round(float(s.min()), 3), "max": round(float(s.max()), 3), "top": s.idxmax(), "bottom": s.idxmin()}


def main():
    rng = np.random.default_rng(SEED)
    sc = pd.read_parquet(PROC / "store_category.parquet"); sp = pd.read_parquet(PROC / "store_panel.parquet").set_index("id")
    base = sc[sc.region_final != "Inconnue"]
    variants = {"baseline_S_R": base, "S_R0_dbfirst_regions": sc[sc.region_db.notna()].assign(region_final=lambda x: x.region_dbfirst).query("region_final != 'Inconnue'"),
                "cert_assoc_sources_only": base[base.first_source.isin(CERT_SOURCES)], "excl_Montreal": base[base.region_final != "Montréal"],
                "excl_top1pct_stores": base[~base.store_id.isin(sc.groupby("store_id").n.sum().pipe(lambda s: s[s > s.quantile(0.99)]).index)],
                "macro8_categories": base.assign(category=base.category.map(MACRO)).groupby(["store_id", "category", "region_final", "region_source", "first_source", "store_kind"], as_index=False).n.sum(),
                "drop_autre": base[base.category != "autre"], "excl_resellers_collectives": base[~base.store_kind.isin(["revendeur", "collectif"])],
                "excl_region_conflicts": base[~base.store_id.map(sp.region_conflict).fillna(False)]}
    res = {}
    for k, d in variants.items():
        res[k] = {"h1_M": h1(d, rng, "M"), "h1_T": h1(d, rng, "T"), "entropy_M": ent_range(d), "gini_listings": round(gini(d.groupby("store_id").n.sum().values), 4)}
        log.info("%s %s", k, json.dumps(res[k]))
    audit = BASE_EXP / "category_audit.json"
    if audit.exists():
        conf = json.load(open(audit)).get("confusion_rule_vs_consensus", {})
        if conf:
            cats = sorted(sc.category.unique()); Vs = []
            P_ = {r: {c: v for c, v in row.items()} for r, row in conf.items()}
            for k in range(100):
                d = base.copy()
                def draw(c):
                    row = P_.get(c); 
                    if not row: return c
                    ks, vs = list(row.keys()), np.array(list(row.values()), float); return rng.choice(ks, p=vs / vs.sum())
                d["category"] = [draw(c) for c in d.category]
                d = d.groupby(["store_id", "category", "region_final", "region_source", "first_source", "store_kind"], as_index=False).n.sum()
                Vs.append(h1(d, rng, "M", nperm=200))
            res["confusion_reassignment_100draws"] = {"V_mean": round(float(np.mean([v["V"] for v in Vs])), 4), "V_min": round(float(min(v["V"] for v in Vs)), 4), "share_p_lt_05": round(float(np.mean([v["p"] < 0.05 for v in Vs])), 3)}
    write_json(res, EXP_DIR / "robustness.json", "robustness"); ledger("validation.robustness", variants=list(res))


if __name__ == "__main__":
    main()