"""Registered test (plan/analysis-plan.md): do bodily-load measures relate to depressive-symptom
severity linearly, or with a break (steeper slope) above a severity threshold?

Data: NHANES 2005-2010 adults with complete PHQ-9 (loader code/nhanes.py, from an earlier
registered analysis). Severity X = PHQ-9 total (0-27).
H1, per outcome (load-bearing limitation, allostatic load, BMI; negative control: height):
  Y = covariates + b1*X + b2*max(0, X - c), c on the integer grid 8..22; statistic
  T = max over c of the t statistic of b2 (one-sided, steeper above c); p from a wild bootstrap
  (Rademacher weights) of the linear-model residuals, recomputing T (Hansen-type sup test).
  Covariates: age, age squared, male, survey cycle. Holm across the three load outcomes.
H2: within strata of PHQ item 10 (functional difficulty, not part of the total): partial Spearman
  of load-bearing limitation with somatic symptoms adjusting for cognitive-affective symptoms
  (as in the earlier analysis), compared between high difficulty (very or extremely difficult)
  and lower difficulty (not at all or somewhat), bootstrap CI of the difference.
Usage: python code/breakpoint.py <nhanes_dir> <out.json>
"""
import json
import sys

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

sys.path.insert(0, "code")
import nhanes as N  # noqa: E402

GRID = list(range(8, 23))
NBOOT = 999
SEED = 20261010
OUTCOMES = {"struct": "load-bearing limitation", "allostatic": "allostatic load", "bmi": "body mass index"}
CONTROL = {"height": "standing height (negative control)"}


def design(d, c=None):
    cyc = pd.get_dummies(d["cycle"], drop_first=True, dtype=float)
    X = np.column_stack([np.ones(len(d)), d["RIDAGEYR"], d["RIDAGEYR"] ** 2, d["male"], cyc.to_numpy(), d["phq"]])
    if c is not None:
        X = np.column_stack([X, np.maximum(0, d["phq"] - c)])
    return X


def sup_t(y, Xs):
    best, best_c, best_b = -np.inf, None, None
    for c, X in Xs.items():
        b, res, *_ = np.linalg.lstsq(X, y, rcond=None)
        r = y - X @ b
        s2 = r @ r / (len(y) - X.shape[1])
        XtXi = np.linalg.inv(X.T @ X)
        t = b[-1] / np.sqrt(s2 * XtXi[-1, -1])
        if t > best:
            best, best_c, best_b = t, c, (b[-2], b[-1])
    return best, best_c, best_b


def test(d, col, rng):
    d = d.dropna(subset=[col]).reset_index(drop=True)
    y = d[col].to_numpy(float)
    X0 = design(d)
    b0, *_ = np.linalg.lstsq(X0, y, rcond=None)
    fit0, r0 = X0 @ b0, y - X0 @ b0
    Xs = {c: design(d, c) for c in GRID}
    T, c_hat, (slope_below, slope_change) = sup_t(y, Xs)
    boot = np.empty(NBOOT)
    for i in range(NBOOT):
        ystar = fit0 + r0 * rng.choice([-1.0, 1.0], len(y))
        boot[i] = sup_t(ystar, Xs)[0]
    sd = float(np.std(y))
    return {"n": int(len(y)), "breakpoint_phq": int(c_hat), "sup_t": round(float(T), 3),
            "p_one_sided": float((np.sum(boot >= T) + 1) / (NBOOT + 1)),
            "slope_below_per_point_sd": round(float(slope_below / sd), 5),
            "slope_change_per_point_sd": round(float(slope_change / sd), 5),
            "linear_slope_per_point_sd": round(float(b0[-1] / sd), 5)}


def pcor(df, x, y, covars):
    def res(v):
        Y = stats.rankdata(df[v]); C = np.column_stack([np.ones(len(df))] + [stats.rankdata(df[c]) for c in covars])
        return Y - C @ np.linalg.lstsq(C, Y, rcond=None)[0]
    return float(np.corrcoef(res(x), res(y))[0, 1])


def holm(ps):
    order = sorted(ps, key=ps.get); out, run = {}, 0.0
    for i, k in enumerate(order):
        run = max(run, min(1.0, (len(ps) - i) * ps[k])); out[k] = round(run, 5)
    return out


def main():
    src, out = sys.argv[1], sys.argv[2]
    rng = np.random.default_rng(SEED)
    d = N.load(src)
    a = N.build(d)
    a["phq"] = a[N.PHQ].sum(axis=1)
    a["height"] = a["BMXHT"]
    a = a.dropna(subset=["RIDAGEYR", "male"])
    res = {"n_adults": int(len(a)), "phq_bands": {f"{lo}_{hi}": int(((a.phq >= lo) & (a.phq <= hi)).sum()) for lo, hi in ((0, 4), (5, 9), (10, 14), (15, 19), (20, 27))}}
    H1 = {k: test(a, k, rng) for k in OUTCOMES}
    for k, v in holm({k: v["p_one_sided"] for k, v in H1.items()}).items():
        H1[k]["p_holm"] = v
    res["H1"] = H1
    res["control_height"] = test(a, "height", rng)
    res["breakpoints"] = {k: v["breakpoint_phq"] for k, v in H1.items()}
    # H2: functional-difficulty strata (DPQ100 asked when any symptom > 0)
    b = a.dropna(subset=["struct", "somatic", "cognitive", "DPQ100"])
    b = b[b["DPQ100"].isin([0, 1, 2, 3])]
    hi, lo = b[b["DPQ100"] >= 2], b[b["DPQ100"] <= 1]
    covs = ["cognitive", "RIDAGEYR", "male"]
    r_hi, r_lo = pcor(hi, "struct", "somatic", covs), pcor(lo, "struct", "somatic", covs)
    diffs = []
    for _ in range(2000):
        h = hi.sample(len(hi), replace=True, random_state=int(rng.integers(1 << 31)))
        l = lo.sample(len(lo), replace=True, random_state=int(rng.integers(1 << 31)))
        diffs.append(pcor(h, "struct", "somatic", covs) - pcor(l, "struct", "somatic", covs))
    res["H2_difficulty_strata"] = {"n_high_difficulty": int(len(hi)), "n_lower_difficulty": int(len(lo)),
                                   "pcor_high": round(r_hi, 4), "pcor_lower": round(r_lo, 4), "difference": round(r_hi - r_lo, 4),
                                   "ci95": [round(float(np.percentile(diffs, 2.5)), 4), round(float(np.percentile(diffs, 97.5)), 4)],
                                   "p_one_sided_greater": float((np.sum(np.array(diffs) <= 0) + 1) / (len(diffs) + 1))}
    with open(out, "w") as fh:
        json.dump(res, fh, indent=2, sort_keys=True); fh.write("\n")


if __name__ == "__main__":
    main()