#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — contact@spboucher.ai
#
"""Inter-quantile Wald tests with a joint bootstrap covariance (October 2026
revision).

The published test (script 02) treats the tau = 0.10 and tau = 0.90 estimates
as independent, i.e. it omits the covariance term of Eq. (iqr_test). Here the
two quantile regressions are re-estimated on the same 150,000-observation
subsample (seed 42) and the difference beta(0.10) - beta(0.90) is bootstrapped
jointly: (i) a paired observation bootstrap and (ii) a paired ZIP3-cluster
bootstrap that also allows for spatial dependence within three-digit ZIP areas.
Each bootstrap draw re-fits both quantile regressions (about 15 s per draw);
draws run in parallel. Default: 200 draws, roughly 8-10 minutes on 16 workers.

Produces results/revision/iqr_wald_bootstrap.csv.

Usage:  python scripts/07_revision_qr_bootstrap.py [--draws 200] [--workers 8]
"""

import argparse
import sys
import time
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path

import numpy as np
import pandas as pd
import statsmodels.api as sm
from scipy import stats

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from wp3 import config, data  # noqa: E402

OUT = config.RESULTS_DIR / "revision"
_X = _Y = _CODES = None  # populated in each worker via initializer


def _init(X, y, codes):
    global _X, _Y, _CODES
    _X, _Y, _CODES = X, y, codes


def _fit_pair(idx):
    r10 = sm.QuantReg(_Y[idx], _X[idx]).fit(q=0.10, max_iter=2000)
    r90 = sm.QuantReg(_Y[idx], _X[idx]).fit(q=0.90, max_iter=2000)
    return r10.params, r90.params


def _draw(args):
    seed, scheme = args
    rng = np.random.RandomState(seed)
    n = len(_Y)
    if scheme == "obs":
        idx = rng.randint(0, n, n)
    else:
        uniq = np.unique(_CODES)
        members = {g: np.where(_CODES == g)[0] for g in uniq}
        draw = rng.choice(uniq, size=len(uniq), replace=True)
        idx = np.concatenate([members[g] for g in draw])
    b10, b90 = _fit_pair(idx)
    return b10 - b90


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--draws", type=int, default=200)
    parser.add_argument("--workers", type=int, default=8)
    args = parser.parse_args()
    OUT.mkdir(parents=True, exist_ok=True)
    t0 = time.time()

    md = data.load_model_data()
    df = md["df_clean"]
    X_const, y = md["X_const"], md["y_clean"]
    rng = np.random.RandomState(config.RANDOM_STATE)
    idx = rng.choice(len(X_const), size=config.QR_SUBSAMPLE_SIZE, replace=False)
    X = X_const.iloc[idx].values.astype(float)
    yy = y.iloc[idx].values.astype(float)
    codes = pd.factorize(df["address_zipcode"].fillna("000").astype(str).str[:3].iloc[idx])[0]
    names = list(X_const.columns)
    kpos = [names.index(v) for v in config.KEY_VARS]

    print("Point estimates on the seed-42 subsample...")
    _init(X, yy, codes)
    r10 = sm.QuantReg(yy, X).fit(q=0.10, max_iter=2000)
    r90 = sm.QuantReg(yy, X).fit(q=0.90, max_iter=2000)
    diff = r10.params - r90.params
    se_indep = np.sqrt(r10.bse ** 2 + r90.bse ** 2)

    results = {}
    for scheme in ["obs", "zip3"]:
        print(f"Bootstrap ({scheme}), {args.draws} draws on {args.workers} workers...")
        t = time.time()
        with ProcessPoolExecutor(max_workers=args.workers, initializer=_init,
                                 initargs=(X, yy, codes)) as ex:
            draws = list(ex.map(_draw, [(1000 + i, scheme) for i in range(args.draws)]))
        results[scheme] = np.vstack(draws)
        print(f"   done in {time.time() - t:.0f}s")

    rows = []
    for v, j in zip(config.KEY_VARS, kpos):
        d_obs = results["obs"][:, j]
        d_z3 = results["zip3"][:, j]
        row = {"variable": v, "beta_010": r10.params[j], "beta_090": r90.params[j],
               "difference": diff[j],
               "se_independent": se_indep[j], "z_independent": diff[j] / se_indep[j],
               "se_boot_obs": d_obs.std(ddof=1), "z_boot_obs": diff[j] / d_obs.std(ddof=1),
               "ci_low_obs": np.percentile(d_obs, 2.5), "ci_high_obs": np.percentile(d_obs, 97.5),
               "se_boot_zip3": d_z3.std(ddof=1), "z_boot_zip3": diff[j] / d_z3.std(ddof=1),
               "ci_low_zip3": np.percentile(d_z3, 2.5), "ci_high_zip3": np.percentile(d_z3, 97.5)}
        row["p_boot_zip3"] = 2 * (1 - stats.norm.cdf(abs(row["z_boot_zip3"])))
        row["significant_5pct_zip3"] = bool(abs(row["z_boot_zip3"]) > 1.96)
        rows.append(row)
    out = pd.DataFrame(rows)
    out["n_draws"] = args.draws
    out["n_zip3_clusters"] = len(np.unique(codes))
    out.to_csv(OUT / "iqr_wald_bootstrap.csv", index=False)
    print(out[["variable", "difference", "z_independent", "z_boot_obs", "z_boot_zip3",
               "significant_5pct_zip3"]].round(3).to_string(index=False))
    print(f"\nDone in {time.time() - t0:.0f}s")


if __name__ == "__main__":
    main()