"""Registered analysis of network timescales and dominance dynamics (see plan/analysis-plan.md).

Usage: python code/analyze.py <features.csv> <participants.tsv> <out.json>
"""
import json
import sys

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

SEED = 20261007
N_BOOT = 10000
QC = {"mean_fd_mm": 0.3, "share_fd_over_0_5": 0.2, "max_translation_mm": 3.0, "registration_dice_min": 0.8}


def auc(a, b):
    a, b = np.asarray(a), np.asarray(b)
    return float((a[:, None] > b[None, :]).mean() + 0.5 * (a[:, None] == b[None, :]).mean())


def boot(fn, arrays, rng):
    v = [fn(*[x[rng.integers(0, len(x), len(x))] for x in arrays]) for _ in range(N_BOOT)]
    return [round(float(q), 4) for q in np.percentile(v, [2.5, 97.5])]


def paired(a, b, rng):
    d = np.asarray(a) - np.asarray(b)
    hl = lambda x: float(np.median(((x[:, None] + x[None, :]) / 2)[np.triu_indices(len(x))]))
    return {"n": int(len(d)), "hodges_lehmann_s": round(hl(d), 4), "ci95": boot(hl, [d], rng),
            "n_positive": int((d > 0).sum()), "p_one_sided": float(stats.wilcoxon(d, alternative="greater").pvalue)}


def groups(dep, ctl, rng):
    dep, ctl = np.asarray(dep, float), np.asarray(ctl, float)
    return {"n_depression": int(len(dep)), "n_control": int(len(ctl)),
            "median_depression": round(float(np.median(dep)), 4), "median_control": round(float(np.median(ctl)), 4),
            "auc_depression_gt_control": round(auc(dep, ctl), 4), "auc_ci95": boot(auc, [dep, ctl], rng),
            "p_two_sided": float(stats.mannwhitneyu(dep, ctl, alternative="two-sided").pvalue)}


def spearman(x, y, rng):
    x, y = np.asarray(x, float), np.asarray(y, float)
    ok = np.isfinite(x) & np.isfinite(y)
    x, y = x[ok], y[ok]
    r = stats.spearmanr(x, y)
    rs = []
    for _ in range(N_BOOT):
        i = rng.integers(0, len(x), len(x))
        rs.append(stats.spearmanr(x[i], y[i]).statistic)
    return {"n": int(len(x)), "rho": round(float(r.statistic), 4),
            "ci95": [round(float(q), 4) for q in np.nanpercentile(rs, [2.5, 97.5])],
            "p_two_sided": float(r.pvalue)}


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


def rank_regression(df, y, covars):
    d = df[[y, "dep"] + covars].dropna()
    Y = stats.rankdata(d[y])
    X = np.column_stack([np.ones(len(d)), d["dep"]] + [stats.rankdata(d[c]) for c in covars])
    b, *_ = np.linalg.lstsq(X, Y, rcond=None)
    r = Y - X @ b
    s2 = r @ r / (len(Y) - X.shape[1])
    se = np.sqrt(np.diag(s2 * np.linalg.inv(X.T @ X)))
    t = b[1] / se[1]
    return {"n": int(len(d)), "t_depression": round(float(t), 3),
            "p_two_sided": float(2 * stats.t.sf(abs(t), len(Y) - X.shape[1]))}


def main():
    rng = np.random.default_rng(SEED)
    f = pd.read_csv(sys.argv[1])
    p = pd.read_csv(sys.argv[2], sep="\t", na_values=["n/a"])
    d = f.merge(p, on="participant_id", how="left")
    res = {"qc_thresholds": QC, "subjects_processed": int(len(d))}
    keep = ((d.mean_fd_mm <= QC["mean_fd_mm"]) & (d.share_fd_over_0_5 <= QC["share_fd_over_0_5"])
            & (d.max_translation_mm <= QC["max_translation_mm"]) & (d.registration_dice >= QC["registration_dice_min"]))
    res["excluded"] = sorted(d.loc[~keep, "participant_id"].tolist())
    d = d[keep].copy()
    d["dep"] = (d["group"] == "depr").astype(float)
    dep, ctl = d[d.group == "depr"], d[d.group == "control"]
    res["sample"] = {"included": int(len(d)), "depression": int(len(dep)), "control": int(len(ctl))}
    # Family A: timescale hierarchy across all included participants
    a = {"A1_DMN_gt_SN": paired(d.acw50_DMN, d.acw50_SN, rng),
         "A2_DMN_gt_CEN": paired(d.acw50_DMN, d.acw50_CEN, rng),
         "A3_CEN_gt_SN": paired(d.acw50_CEN, d.acw50_SN, rng)}
    for k, v in zip(a, holm([v["p_one_sided"] for v in a.values()])):
        a[k]["p_holm"] = v
    res["family_A_hierarchy_acw50"] = a
    # Family B: depression
    b = {"B1_acw50_DMN_group": groups(dep.acw50_DMN, ctl.acw50_DMN, rng),
         "B2_dmn_dwell_group": groups(dep.dmn_dwell_s, ctl.dmn_dwell_s, rng),
         "B3_transition_entropy_group": groups(dep.transition_entropy_bits, ctl.transition_entropy_bits, rng),
         "B4_brooding_vs_acw50_DMN": spearman(dep["RRS-brooding"], dep.acw50_DMN, rng),
         "B5_brooding_vs_dmn_dwell": spearman(dep["RRS-brooding"], dep.dmn_dwell_s, rng)}
    for k, v in zip(b, holm([v["p_two_sided"] for v in b.values()])):
        b[k]["p_holm"] = v
    res["family_B_depression"] = b
    # Secondary (unadjusted)
    s = {}
    for net in ("DMN", "SN", "CEN"):
        s[f"S_acw0_{net}_group"] = groups(dep[f"acw0_{net}"], ctl[f"acw0_{net}"], rng)
    for net in ("SN", "CEN"):
        s[f"S_acw50_{net}_group"] = groups(dep[f"acw50_{net}"], ctl[f"acw50_{net}"], rng)
    s["S_acw50_all_group"] = groups(dep.acw50_all, ctl.acw50_all, rng)
    s["S_dmn_occupancy_group"] = groups(dep.dmn_occupancy, ctl.dmn_occupancy, rng)
    s["S_BDI_vs_acw50_DMN_all"] = spearman(d.BDI, d.acw50_DMN, rng)
    s["S_Zung_vs_acw50_DMN_all"] = spearman(d.Zung_SDS, d.acw50_DMN, rng)
    s["S_RRS_sum_vs_acw50_DMN_dep"] = spearman(dep["RRS-sum"], dep.acw50_DMN, rng)
    s["S_mean_fd_group"] = groups(dep.mean_fd_mm, ctl.mean_fd_mm, rng)
    s["S_motion_age_sex_adjusted_acw50_DMN"] = rank_regression(d.assign(male=(d.gender.str.lower() == "m").astype(float)),
                                                               "acw50_DMN", ["mean_fd_mm", "age", "male"])
    s["S_motion_age_sex_adjusted_dmn_dwell"] = rank_regression(d.assign(male=(d.gender.str.lower() == "m").astype(float)),
                                                               "dmn_dwell_s", ["mean_fd_mm", "age", "male"])
    res["secondary"] = s
    with open(sys.argv[3], "w") as fh:
        json.dump(res, fh, indent=2, sort_keys=True)
        fh.write("\n")
    print(json.dumps(res, indent=1)[:4000])


if __name__ == "__main__":
    main()