"""H2 — Regional assortment diversity vs regional scale (n <= 17; preregistered: Spearman rho with CI; OLS HC3 secondary),
plus H7-V3 (incremental information of the ECI beyond population, business count and income; leave-one-region-out).
External covariates: ISQ population 2025p, ISQ business locations with employees Dec 2025, ISQ disposable income per capita 2024p,
Census 2021 land area (density) — see data/external/README.md. Output experiments/h2_scale.json"""
from __future__ import annotations
import json
import numpy as np, pandas as pd
from scipy import stats
from src.common import ROOT, SEED, get_logger, ledger, write_json, PROC_DIR, EXP_DIR, P

log = get_logger("models.h2")
PROC = PROC_DIR
COV = {"log_population": ("population", "ISQ population estimate, 1 July 2025 (provisional)"), "log_business_locations": ("business_locations_with_employees", "ISQ business locations with employees, Dec 2025"),
       "disposable_income": ("disposable_income_per_capita_cad", "ISQ disposable income per capita, 2024p"), "log_density": ("population_density_per_km2", "Census 2021 population density (2021 ER boundaries)")}


def spearman_ci(x, y, rng, B=5000, boot_y=None):
    """Region-level bootstrap combined with outcome resampling draws (boot_y: array B x n of outcome draws, optional)."""
    n = len(x); rhos = []
    for b in range(B):
        idx = rng.integers(0, n, n); yy = y if boot_y is None else boot_y[b % boot_y.shape[0]]
        if len(set(idx)) < 4: continue
        rhos.append(stats.spearmanr(x[idx], yy[idx]).correlation)
    rhos = np.array(rhos); return float(stats.spearmanr(x, y).correlation), [float(np.nanpercentile(rhos, 2.5)), float(np.nanpercentile(rhos, 97.5))]


def ols_hc3(x, y):
    X = np.column_stack([np.ones_like(x), x]); b = np.linalg.lstsq(X, y, rcond=None)[0]; e = y - X @ b; h = np.diag(X @ np.linalg.pinv(X.T @ X) @ X.T)
    V = np.linalg.pinv(X.T @ X) @ (X.T * (e ** 2 / (1 - h) ** 2)) @ X @ np.linalg.pinv(X.T @ X); return float(b[1]), float(np.sqrt(V[1, 1]))


def main():
    rng = np.random.default_rng(SEED)
    ext = pd.read_csv(ROOT / "data" / "external" / "regions_external.csv").set_index("region_name")
    div = pd.read_csv(PROC / "diversity_region.csv").set_index("region"); ind = pd.read_csv(PROC / "h7_region_indicators.csv").set_index("region")
    sp = pd.read_parquet(PROC / "store_panel.parquet"); cnt = sp[sp.region_final != "Inconnue"].groupby("region_final").size()
    df = div.join(ext, how="inner"); df["n_stores"] = cnt.reindex(df.index)
    for k, (col, _) in COV.items():
        df[k] = np.log(df[col]) if k.startswith("log") else df[col]
    out = {"regions_in_external": int(len(ext)), "regions_matched": int(len(df)), "correlations": {}, "covariate_sources": {k: v[1] for k, v in COV.items()},
           "note": "17 regions at most (Nord-du-Québec has no located store); rarefied diversity needs >= 15 stores. Associations only; CIs from 5 000 region-level bootstrap draws; OLS slope with HC3 SE."}
    for outc in ("H_rarefied15", "H_norm_M", "eff_cat_M"):
        for k in COV:
            d = df[[outc, k]].dropna(); x = d[k].values.astype(float); y = d[outc].values.astype(float)
            if len(d) < 6: continue
            rho, ci = spearman_ci(x, y, rng); b, se = ols_hc3(x, y)
            out["correlations"][f"{outc} ~ {k}"] = {"n": int(len(d)), "rho": round(rho, 3), "ci": [round(ci[0], 3), round(ci[1], 3)], "ols_b": round(b, 4), "ols_se_hc3": round(se, 4)}
    # stores per capita vs scale (coverage gradient)
    df["log_stores_per_100k"] = np.log(df.n_stores / df.population * 1e5)
    for k in ("log_population", "log_density", "disposable_income"):
        d = df[["log_stores_per_100k", k]].dropna(); rho, ci = spearman_ci(d[k].values.astype(float), d.log_stores_per_100k.values, rng)
        out["correlations"][f"log_stores_per_100k ~ {k}"] = {"n": int(len(d)), "rho": round(rho, 3), "ci": [round(ci[0], 3), round(ci[1], 3)], "ols_b": None, "ols_se_hc3": None}
    # H7-V3: ECI on covariates, leave-one-region-out predictive correlation
    e = ind.join(df[["log_population", "log_business_locations", "disposable_income"]], how="inner").dropna(subset=["eci"])
    X = np.column_stack([np.ones(len(e)), e.log_population, e.log_business_locations, e.disposable_income]); y = e.eci.values
    b = np.linalg.lstsq(X, y, rcond=None)[0]; r2 = 1 - ((y - X @ b) ** 2).sum() / ((y - y.mean()) ** 2).sum()
    pred = np.array([(X[i] @ np.linalg.lstsq(np.delete(X, i, 0), np.delete(y, i), rcond=None)[0]) for i in range(len(y))])
    v3 = {"n": int(len(y)), "R2_eci_on_covariates": round(float(r2), 3), "loo_pred_spearman": round(float(stats.spearmanr(pred, y).correlation), 3), "criterion_rho_ge_0.5": bool(stats.spearmanr(pred, y).correlation >= 0.5)}
    xd = e.join(ind[["diversity"]], rsuffix="_i"); v3["partial_corr_eci_diversity_given_covariates"] = round(float(stats.spearmanr(y - X @ b, xd.diversity - X @ np.linalg.lstsq(X, xd.diversity.values, rcond=None)[0]).correlation), 3)
    out["H7_V3"] = v3
    df.reset_index().rename(columns={"index": "region"}).to_csv(PROC / "h2_region_table.csv", index=False)
    write_json(out, EXP_DIR / "h2_scale.json", "h2"); ledger("models.h2_scale", **{k: v for k, v in out.items() if k != "correlations"}, primary=out["correlations"].get("H_rarefied15 ~ log_population"))
    log.info(json.dumps({k: v for k, v in out["correlations"].items() if "H_rarefied15" in k or "per_100k" in k})); log.info(json.dumps(v3))


if __name__ == "__main__":
    main()