"""Registered analysis (see plan/analysis-plan.md): do measures of bodily load, structural
limitation and slow bowel transit relate to specific somatic depressive symptoms, and to
somatic more than to cognitive-affective symptoms, in NHANES 2005-2010 adults?

Usage: python code/analyze.py <dir with NHANES .xpt files> <out.json>
"""
import json
import sys

import numpy as np
import pandas as pd
from scipy import stats

SEED = 20261007
N_BOOT = int(__import__("os").environ.get("N_BOOT", 2000))
CYCLES = "DEF"
FILES = ["DPQ", "BHQ", "PFQ", "BMX", "BPX", "CRP", "GHB", "TCHOL", "HDL", "BIOPRO"]
PHQ = [f"DPQ0{i}0" for i in range(1, 10)]
SOMATIC = ["DPQ030", "DPQ040", "DPQ050", "DPQ080"]          # sleep, energy, appetite, psychomotor
COGNITIVE = ["DPQ010", "DPQ020", "DPQ060", "DPQ070", "DPQ090"]  # anhedonia, mood, worth, concentration, suicidality
VEGETATIVE = ["DPQ030", "DPQ040", "DPQ050"]
STRUCT_ITEMS = ["PFQ061B", "PFQ061C", "PFQ061D", "PFQ061E", "PFQ061I", "PFQ061M", "PFQ061T"]
SCREENERS = ["PFQ049", "PFQ051", "PFQ054", "PFQ057", "PFQ059"]


def load(src):
    parts = []
    for s in CYCLES:
        d = pd.read_sas(f"{src}/DEMO_{s}.xpt", format="xport")[["SEQN", "RIDAGEYR", "RIAGENDR", "INDFMPIR"]]
        for f in FILES:
            d = d.merge(pd.read_sas(f"{src}/{f}_{s}.xpt", format="xport"), on="SEQN", how="left")
        d["cycle"] = s
        parts.append(d)
    d = pd.concat(parts, ignore_index=True)
    num = d.select_dtypes("number").columns
    d[num] = d[num].round(6)  # SAS transport zeros are read as tiny non-zero floats
    return d


def structural(row):
    vals = {k: row[k] for k in STRUCT_ITEMS}
    if row["PFQ054"] == 1:  # needs equipment to walk: walking items not asked, scored unable
        for k in ("PFQ061B", "PFQ061C"):
            if not (1 <= (vals[k] or 0) <= 4):
                vals[k] = 4
    asked = [v for v in vals.values() if v in (1, 2, 3, 4)]
    if len(asked) >= 5:
        return float(np.mean([v - 1 for v in asked]))
    screened_out = row["RIDAGEYR"] < 60 and all(row[s] == 2 for s in SCREENERS if not pd.isna(row[s])) \
        and not pd.isna(row["PFQ049"]) and all(pd.isna(v) for v in vals.values())
    return 0.0 if screened_out else np.nan


def build(d):
    a = d[d.RIDAGEYR >= 20].copy()
    ok = a[PHQ].isin([0, 1, 2, 3]).all(axis=1)
    a = a[ok].copy()
    a["somatic"] = a[SOMATIC].mean(axis=1)
    a["cognitive"] = a[COGNITIVE].mean(axis=1)
    a["vegetative"] = a[VEGETATIVE].mean(axis=1)
    a["psychomotor"] = a["DPQ080"]
    a["struct"] = a.apply(structural, axis=1)
    bm = a["BHD050"].where(a["BHD050"].between(0, 200))
    stool = a["BHQ060"].where(a["BHQ060"].between(1, 7))
    # slow transit: Bristol type 1-2 or fewer than 3 bowel movements a week; if one answer is
    # missing, slow when the other indicates slow, otherwise missing
    slow_s = stool <= 2
    slow_b = bm < 3
    known_both = stool.notna() & bm.notna()
    slow_any = (slow_s & stool.notna()) | (slow_b & bm.notna())
    a["slow_transit"] = np.where(slow_any, 1.0, np.where(known_both, 0.0, np.nan))
    a["bmi"] = a["BMXBMI"]
    a["sbp"] = a[["BPXSY1", "BPXSY2", "BPXSY3"]].mean(axis=1)
    a["dbp"] = a[["BPXDI1", "BPXDI2", "BPXDI3"]].replace(0, np.nan).mean(axis=1)
    markers = {"sbp": "high", "dbp": "high", "BMXWAIST": "high", "LBXGH": "high", "LBXCRP": "high",
               "LBXTC": "high", "LBDHDD": "low", "LBXSAL": "low", "LBXSCR": "high"}
    flags = []
    for m, side in markers.items():
        q = a[m].quantile(0.75 if side == "high" else 0.25)
        f = (a[m] >= q) if side == "high" else (a[m] <= q)
        flags.append(np.where(a[m].notna(), f.astype(float), np.nan))
    F = np.column_stack(flags)
    n_avail = np.isfinite(F).sum(axis=1)
    with np.errstate(invalid="ignore", divide="ignore"):
        a["allostatic"] = np.where(n_avail >= 7, np.nansum(F, axis=1) / n_avail * len(markers), np.nan)
    a["low_income"] = -a["INDFMPIR"]
    a["male"] = (a["RIAGENDR"] == 1).astype(float)
    return a


def residual_ranks(df, y, covars):
    Y = stats.rankdata(df[y])
    X = np.column_stack([np.ones(len(df))] + [stats.rankdata(df[c]) if df[c].nunique() > 2 else df[c].to_numpy(float)
                                               for c in covars])
    b, *_ = np.linalg.lstsq(X, Y, rcond=None)
    return Y - X @ b


def pcor(df, x, y, covars):
    rx, ry = residual_ranks(df, x, covars), residual_ranks(df, y, covars)
    return float(np.corrcoef(rx, ry)[0, 1])


def boot(df, fn, rng):
    n = len(df)
    vals = []
    for _ in range(N_BOOT):
        vals.append(fn(df.iloc[rng.integers(0, n, n)]))
    return [round(float(v), 4) for v in np.percentile(vals, [2.5, 97.5])]


def holm(ps):
    order = np.argsort(ps)
    adj, run = [0.0] * len(ps), 0.0
    for k, i in enumerate(order):
        run = max(run, min(1.0, ps[i] * (len(ps) - k)))
        adj[i] = round(run, 5)
    return adj


def boot_p(df, fn, rng):
    """Two-sided bootstrap p for a statistic being zero (percentile method)."""
    n = len(df)
    vals = np.array([fn(df.iloc[rng.integers(0, n, n)]) for _ in range(N_BOOT)])
    p = 2 * min((vals <= 0).mean(), (vals >= 0).mean())
    return max(p, 1 / N_BOOT), [round(float(v), 4) for v in np.percentile(vals, [2.5, 97.5])]


def main():
    rng = np.random.default_rng(SEED)
    a = build(load(sys.argv[1]))
    if "--shuffle-phq" in sys.argv:  # pipeline test only: break every PHQ association
        perm = np.random.default_rng(1).permutation(len(a))
        for c in PHQ + ["somatic", "cognitive", "vegetative", "psychomotor"]:
            a[c] = a[c].to_numpy()[perm]
    base = ["RIDAGEYR", "male", "cycle_E", "cycle_F"]
    a["cycle_E"] = (a.cycle == "E").astype(float)
    a["cycle_F"] = (a.cycle == "F").astype(float)
    res = {"adults_with_phq9": int(len(a)),
           "depressed_phq9_ge_10": int((a[PHQ].sum(axis=1) >= 10).sum()),
           "missing": {c: int(a[c].isna().sum()) for c in ["struct", "slow_transit", "bmi", "allostatic", "low_income"]}}
    # H1: double dissociation
    h = a.dropna(subset=["struct", "slow_transit"])
    cov = base
    d1 = lambda x: pcor(x, "struct", "psychomotor", cov + ["slow_transit"]) - pcor(x, "slow_transit", "psychomotor", cov + ["struct"])
    d2 = lambda x: pcor(x, "slow_transit", "vegetative", cov + ["struct"]) - pcor(x, "struct", "vegetative", cov + ["slow_transit"])
    p1, ci1 = boot_p(h, d1, rng)
    p2, ci2 = boot_p(h, d2, rng)
    res["H1"] = {"n": int(len(h)),
                 "struct_psychomotor": round(pcor(h, "struct", "psychomotor", cov + ["slow_transit"]), 4),
                 "transit_psychomotor": round(pcor(h, "slow_transit", "psychomotor", cov + ["struct"]), 4),
                 "transit_vegetative": round(pcor(h, "slow_transit", "vegetative", cov + ["struct"]), 4),
                 "struct_vegetative": round(pcor(h, "struct", "vegetative", cov + ["slow_transit"]), 4),
                 "H1a_difference": round(d1(h), 4), "H1a_ci95": ci1, "H1a_p": p1,
                 "H1b_difference": round(d2(h), 4), "H1b_ci95": ci2, "H1b_p": p2}
    # H2: somatic beyond cognitive-affective, and specificity
    h2 = {}
    for m in ["struct", "slow_transit", "bmi", "allostatic"]:
        s = a.dropna(subset=[m])
        som = lambda x, m=m: pcor(x, m, "somatic", base + ["cognitive"])
        cog = lambda x, m=m: pcor(x, m, "cognitive", base + ["somatic"])
        p_s, ci_s = boot_p(s, som, rng)
        p_d, ci_d = boot_p(s, lambda x: som(x) - cog(x), rng)
        h2[m] = {"n": int(len(s)), "somatic_adj_cognitive": round(som(s), 4), "somatic_ci95": ci_s, "somatic_p": p_s,
                 "cognitive_adj_somatic": round(cog(s), 4), "difference": round(som(s) - cog(s), 4),
                 "difference_ci95": ci_d, "difference_p": p_d}
    for k, v in zip(h2, holm([v["somatic_p"] for v in h2.values()])):
        h2[k]["somatic_p_holm"] = v
    for k, v in zip(h2, holm([v["difference_p"] for v in h2.values()])):
        h2[k]["difference_p_holm"] = v
    res["H2"] = h2
    # H3: low income
    s = a.dropna(subset=["low_income"])
    som = lambda x: pcor(x, "low_income", "somatic", base + ["cognitive"])
    cog = lambda x: pcor(x, "low_income", "cognitive", base + ["somatic"])
    p3, ci3 = boot_p(s, lambda x: som(x) - cog(x), rng)
    res["H3"] = {"n": int(len(s)), "somatic_adj_cognitive": round(som(s), 4), "cognitive_adj_somatic": round(cog(s), 4),
                 "difference": round(som(s) - cog(s), 4), "difference_ci95": ci3, "difference_p": p3}
    # Secondary: each PHQ item's partial association with each measure (adjusted for base covariates)
    sec = {}
    for m in ["struct", "slow_transit", "bmi", "allostatic", "low_income"]:
        s = a.dropna(subset=[m])
        sec[m] = {it: round(pcor(s, m, it, base), 4) for it in PHQ}
    res["secondary_item_correlations"] = sec
    with open(sys.argv[2], "w") as fh:
        json.dump(res, fh, indent=2, sort_keys=True)
        fh.write("\n")
    print(json.dumps({k: v for k, v in res.items() if k != "secondary_item_correlations"}, indent=1))


if __name__ == "__main__":
    main()