"""Stratified jackknife sensitivity to NHANES clustering/weights.

Reuses the verifier's independent reconstruction, not the bundle's estimator.
Scores and full-sample allostatic cutpoints are held fixed. Replicate ranks and
correlations are re-estimated. These are sensitivity results, not new claims.
"""
from pathlib import Path
import contextlib
import io
import json
import runpy
import sys
import numpy as np
import pandas as pd
from scipy.stats import t

src = Path(sys.argv[1])
with contextlib.redirect_stdout(io.StringIO()):
    independent = runpy.run_path(str(Path(__file__).with_name('independent_points.py')))
d = independent['d']
demo = pd.concat([pd.read_sas(src / f'DEMO_{s}.xpt', format='xport')
    [['SEQN', 'SDMVSTRA', 'SDMVPSU', 'WTMEC2YR']].round(6) for s in 'DEF'])
assert demo.SEQN.is_unique
pairs = demo.groupby('SDMVSTRA').SDMVPSU.unique().sort_index()
assert all(len(p) >= 2 for p in pairs)
d = d.merge(demo, on='SEQN', validate='one_to_one')
strata = {s: i for i, s in enumerate(pairs.index)}
d['stratum_index'] = d.SDMVSTRA.map(strata)
H = len(strata)
R = sum(len(p) for p in pairs)
df = R - H
base = ['RIDAGEYR', 'male', 'E', 'F']

def weighted_midrank(x, w):
    _, inv = np.unique(x, return_inverse=True)
    mass = np.bincount(inv, weights=w)
    return (np.cumsum(mass) - mass / 2)[inv]

def precision(data, cols, weights):
    good = data[cols].notna().all(axis=1).to_numpy() & (weights > 0)
    w = weights[good]
    values = data.loc[good, cols].to_numpy(float)
    ranks = np.column_stack([weighted_midrank(values[:, i], w) for i in range(values.shape[1])])
    centered = ranks - np.average(ranks, axis=0, weights=w)
    cov = centered.T @ (centered * w[:, None]) / w.sum()
    scale = np.sqrt(cov.diagonal())
    return np.linalg.inv(cov / np.outer(scale, scale))

def pc(P, a, b):
    return float(-P[a, b] / np.sqrt(P[a, a] * P[b, b]))

def estimate(weights):
    result = {}
    complete = d[['struct', 'slow_transit']].notna().all(axis=1).to_numpy()
    for outcome in ['psychomotor', 'vegetative']:
        P = precision(d, ['struct', 'slow_transit', outcome] + base, weights * complete)
        result['H1.struct_' + outcome] = pc(P, 0, 2)
        result['H1.transit_' + outcome] = pc(P, 1, 2)
    result['H1.H1a_difference'] = result['H1.struct_psychomotor'] - result['H1.transit_psychomotor']
    result['H1.H1b_difference'] = result['H1.transit_vegetative'] - result['H1.struct_vegetative']
    for measure in ['struct', 'slow_transit', 'bmi', 'allostatic', 'low_income']:
        P = precision(d, [measure, 'somatic', 'cognitive'] + base, weights)
        prefix = 'H3' if measure == 'low_income' else 'H2.' + measure
        result[prefix + '.somatic_adj_cognitive'] = pc(P, 0, 1)
        result[prefix + '.cognitive_adj_somatic'] = pc(P, 0, 2)
        result[prefix + '.difference'] = pc(P, 0, 1) - pc(P, 0, 2)
    return result

def holm(ps):
    out = {}
    running = 0.
    for i, (name, p) in enumerate(sorted(ps.items(), key=lambda item: item[1])):
        running = max(running, min(1., p * (len(ps) - i)))
        out[name] = running
    return out

out = {'strata': H, 'psus': R, 'replicates': R, 'df': df,
       'method': 'Stratified delete-one-PSU jackknife; donor-stratum remaining weights multiplied by m/(m-1); variance sum((m-1)/m*(replicate-full)^2); t intervals with PSUs-strata df; fixed derived scores/cutpoints; re-estimated ranks.'}
for weighted in [False, True]:
    weights = d.WTMEC2YR.to_numpy() / 3 if weighted else np.ones(len(d))
    point = estimate(weights)
    draws = []
    factors = []
    for stratum, units in pairs.items():
        donor = d.SDMVSTRA.eq(stratum).to_numpy()
        factor = (len(units) - 1) / len(units)
        for unit in units:
            adjustment = np.ones(len(d))
            adjustment[donor] = 1 / factor
            adjustment[donor & d.SDMVPSU.eq(unit).to_numpy()] = 0
            draws.append(estimate(weights * adjustment))
            factors.append(factor)
    summary = {}
    for key, value in point.items():
        se = float(np.sqrt(np.sum([f * (r[key] - value) ** 2 for f, r in zip(factors, draws)])))
        critical = float(t.ppf(.975, df))
        summary[key] = {'estimate': value, 'se': se, 'ci95': [value - critical * se, value + critical * se],
                        'p': float(2 * t.sf(abs(value / se), df))}
    for suffix in ['difference', 'somatic_adj_cognitive']:
        family = {key: value['p'] for key, value in summary.items() if key.startswith('H2.') and key.endswith('.' + suffix)}
        for key, value in holm(family).items(): summary[key]['p_holm'] = value
    out['weighted' if weighted else 'unweighted'] = summary
    if not weighted:
        for key, value in point.items():
            expected = independent['r']
            for name in key.split('.'): expected = expected[name]
            assert abs(value - expected) < 1e-10
print(json.dumps(out, indent=2))