"""Independent point-estimate reconstruction from public NHANES files.

Does not import the submitted program. Uses precision matrices of ranked
variables instead of residualizing two separate regressions per correlation.
Bootstrap intervals remain the responsibility of the full submitted rerun.
"""
from pathlib import Path
import json
import sys
import numpy as np
import pandas as pd

src = Path(sys.argv[1])
columns = {
    'DEMO': ['RIDAGEYR', 'RIAGENDR', 'INDFMPIR'],
    'DPQ': [f'DPQ0{i}0' for i in range(1, 10)],
    'BHQ': ['BHD050', 'BHQ060'],
    'PFQ': ['PFQ049', 'PFQ051', 'PFQ054', 'PFQ057', 'PFQ059'] +
           ['PFQ061' + c for c in ['B', 'C', 'D', 'E', 'I', 'M', 'T']],
    'BMX': ['BMXBMI', 'BMXWAIST'],
    'BPX': ['BPXSY1', 'BPXSY2', 'BPXSY3', 'BPXDI1', 'BPXDI2', 'BPXDI3'],
    'CRP': ['LBXCRP'], 'GHB': ['LBXGH'], 'TCHOL': ['LBXTC'],
    'HDL': ['LBDHDD'], 'BIOPRO': ['LBXSAL', 'LBXSCR']
}
cycles = []
unique_counts = {}
for cycle in ['D', 'E', 'F']:
    data = None
    for module, names in columns.items():
        path = src / f'{module}_{cycle}.xpt'
        table = pd.read_sas(path, format='xport')[['SEQN'] + names].round(6)
        assert table.SEQN.is_unique
        unique_counts[path.name] = len(table)
        data = table if data is None else data.merge(table, on='SEQN', how='left', validate='one_to_one')
    data['E'] = int(cycle == 'E')
    data['F'] = int(cycle == 'F')
    cycles.append(data)
d = pd.concat(cycles, ignore_index=True)
phq = columns['DPQ']
d = d.loc[(d.RIDAGEYR >= 20) & d[phq].isin(range(4)).all(axis=1)].copy()
d['male'] = d.RIAGENDR.eq(1).astype(int)
d['somatic'] = d[['DPQ030', 'DPQ040', 'DPQ050', 'DPQ080']].mean(axis=1)
d['cognitive'] = d[['DPQ010', 'DPQ020', 'DPQ060', 'DPQ070', 'DPQ090']].mean(axis=1)
d['vegetative'] = d[['DPQ030', 'DPQ040', 'DPQ050']].mean(axis=1)
d['psychomotor'] = d.DPQ080
items = d[columns['PFQ'][5:]].copy()
for c in ['PFQ061B', 'PFQ061C']:
    items.loc[d.PFQ054.eq(1) & ~items[c].between(1, 4), c] = 4
valid = items.where(items.isin([1, 2, 3, 4]))
d['struct'] = (valid.mean(axis=1) - 1).where(valid.notna().sum(axis=1) >= 5)
screeners = d[columns['PFQ'][:5]]
screened = (d.RIDAGEYR < 60) & (screeners.isna() | screeners.eq(2)).all(axis=1) & d.PFQ049.notna() & items.isna().all(axis=1)
d.loc[screened & d.struct.isna(), 'struct'] = 0
bm = d.BHD050.where(d.BHD050.between(0, 200))
stool = d.BHQ060.where(d.BHQ060.between(1, 7))
d['slow_transit'] = np.where(stool.le(2) | bm.lt(3), 1., np.where(stool.notna() & bm.notna(), 0., np.nan))
d['bmi'] = d.BMXBMI
d['low_income'] = -d.INDFMPIR
d['sbp'] = d[['BPXSY1', 'BPXSY2', 'BPXSY3']].mean(axis=1)
d['dbp'] = d[['BPXDI1', 'BPXDI2', 'BPXDI3']].replace(0, np.nan).mean(axis=1)
adverse = {}
cutpoints = {}
for c in ['sbp', 'dbp', 'BMXWAIST', 'LBXGH', 'LBXCRP', 'LBXTC', 'LBXSCR', 'LBDHDD', 'LBXSAL']:
    low = c in ['LBDHDD', 'LBXSAL']
    cut = d[c].quantile(.25 if low else .75)
    cutpoints[c] = float(cut)
    flag = d[c].le(cut) if low else d[c].ge(cut)
    adverse[c] = flag.astype(float).where(d[c].notna())
flags = pd.DataFrame(adverse)
d['allostatic'] = (9 * flags.mean(axis=1)).where(flags.notna().sum(axis=1) >= 7)
base = ['RIDAGEYR', 'male', 'E', 'F']
def partial(frame, x, y, extra=()):
    table = frame[[x, y] + base + list(extra)].dropna().rank(method='average')
    inv = np.linalg.inv(table.corr().to_numpy())
    return float(-inv[0, 1] / np.sqrt(inv[0, 0] * inv[1, 1]))

r = {'adults_with_phq9': len(d), 'depressed_phq9_ge_10': int(d[phq].sum(axis=1).ge(10).sum()),
     'missing': {c: int(d[c].isna().sum()) for c in ['struct', 'slow_transit', 'bmi', 'allostatic', 'low_income']}}
h = d.dropna(subset=['struct', 'slow_transit'])
r['H1'] = {'n': len(h),
    'struct_psychomotor': partial(h, 'struct', 'psychomotor', ['slow_transit']),
    'transit_psychomotor': partial(h, 'slow_transit', 'psychomotor', ['struct']),
    'struct_vegetative': partial(h, 'struct', 'vegetative', ['slow_transit']),
    'transit_vegetative': partial(h, 'slow_transit', 'vegetative', ['struct'])}
r['H1']['H1a_difference'] = r['H1']['struct_psychomotor'] - r['H1']['transit_psychomotor']
r['H1']['H1b_difference'] = r['H1']['transit_vegetative'] - r['H1']['struct_vegetative']
r['H2'] = {}
for measure in ['struct', 'slow_transit', 'bmi', 'allostatic', 'low_income']:
    frame = d.dropna(subset=[measure])
    som = partial(frame, measure, 'somatic', ['cognitive'])
    cog = partial(frame, measure, 'cognitive', ['somatic'])
    out = {'n': len(frame), 'somatic_adj_cognitive': som, 'cognitive_adj_somatic': cog, 'difference': som - cog}
    if measure == 'low_income': r['H3'] = out
    else: r['H2'][measure] = out
r['secondary_item_correlations'] = {measure: {item: partial(d, measure, item) for item in phq}
    for measure in ['struct', 'slow_transit', 'bmi', 'allostatic', 'low_income']}
print(json.dumps({'results': r, 'source_unique_rows': unique_counts,
                  'allostatic_cutpoints': cutpoints, 'screened_structural_zeros': int(screened.sum())}, indent=2))