"""Registered test (plan/analysis-plan.md): is the anterior insula a selective meeting point of
time perception, interoception and vestibular processing in the neuroimaging literature, and
of somatic and mood terms of depression?

Data: NeuroQuery v1 (13,459 full-text articles: reported MNI coordinates and body-text term
counts) and the Harvard-Oxford cortical atlas (insula, max-probability, 25% threshold, 2 mm).

A study "activates" a region if any of its reported coordinates lies within RADIUS mm of a voxel
of that region. A study "uses" a term if the term occurs at least K times in its body text.
Association of a term with a region: two-proportion z comparing the activation rate of studies
using the term with that of all other studies.

Usage: python code/convergence.py <out.json> [K]
"""
import json
import sys

import nibabel as nib
import numpy as np
import pandas as pd
import scipy.sparse as sp
from scipy.spatial import cKDTree

RADIUS = 10.0
K_DEFAULT = 2
SEED = 20261007
N_NULL = 10000
POOL = (30, 1000)  # random-term pool: terms used by 30-1000 studies (at the same K)
NQ = "data/neuroquery"
HO = "data/HarvardOxford-cort-maxprob-thr25-2mm.nii.gz"
INSULA_LABEL = 2  # "Insular Cortex" in the Harvard-Oxford cortical atlas (1-based index)
TARGETS = {"time perception": ["time perception"], "interoception": ["interoception", "interoceptive"],
           "vestibular": ["vestibular"]}
SYMPTOMS = {"fatigue": ["fatigue"], "psychomotor": ["psychomotor"], "sadness": ["sad", "sadness"],
            "mood": ["mood"]}


def regions():
    img = nib.load(HO)
    d = np.asarray(img.dataobj)
    ijk = np.argwhere(d == INSULA_LABEL)
    xyz = nib.affines.apply_affine(img.affine, ijk)
    return {"anterior_insula": xyz[xyz[:, 1] >= 0], "posterior_insula": xyz[xyz[:, 1] < 0]}


def main():
    out_path = sys.argv[1]
    k = int(sys.argv[2]) if len(sys.argv) > 2 else K_DEFAULT
    rng = np.random.default_rng(SEED)
    vocab = [l.rstrip("\n").split("\t")[0] for l in open(f"{NQ}/data-neuroquery_version-1_vocab-neuroquery7547_vocabulary.txt")]
    X = sp.load_npz(f"{NQ}/data-neuroquery_version-1_vocab-neuroquery7547_source-body_type-count_features.npz").tocsc()
    meta = pd.read_csv(f"{NQ}/data-neuroquery_version-1_metadata.tsv.gz", sep="\t")
    co = pd.read_csv(f"{NQ}/data-neuroquery_version-1_coordinates.tsv.gz", sep="\t")
    ids = meta["id"].to_numpy()
    row_of = {i: r for r, i in enumerate(ids)}
    co = co[co["id"].isin(row_of)]
    study_row = co["id"].map(row_of).to_numpy()
    pts = co[["x", "y", "z"]].to_numpy(float)
    regs = regions()
    active = {}
    for name, vox in regs.items():
        d, _ = cKDTree(vox).query(pts, k=1)
        hit = np.zeros(len(ids), bool)
        hit[np.unique(study_row[d <= RADIUS])] = True
        active[name] = hit
    has_coords = np.zeros(len(ids), bool)
    has_coords[np.unique(study_row)] = True

    U = (X >= k).astype(np.int8).tocsc()
    U = U[has_coords]
    act = {n: a[has_coords] for n, a in active.items()}
    N = U.shape[0]
    n_term = np.asarray(U.sum(0)).ravel()

    def z_for(cols, region):
        a = act[region]
        m = np.zeros(N, bool)
        for c in cols:
            m |= U[:, c].toarray().ravel().astype(bool)
        n1 = m.sum()
        n0 = N - n1
        p1 = a[m].mean()
        p0 = a[~m].mean()
        p = a.mean()
        z = (p1 - p0) / np.sqrt(p * (1 - p) * (1 / n1 + 1 / n0))
        return float(z), int(n1), float(p1), float(p0)

    # all single terms' z, for the random-term null
    a_all = {}
    for region in act:
        a = act[region].astype(float)
        hits_with = np.asarray(U.T @ a).ravel()
        p = a.mean()
        p1 = np.divide(hits_with, n_term, out=np.zeros_like(hits_with), where=n_term > 0)
        p0 = (a.sum() - hits_with) / (N - n_term)
        with np.errstate(divide="ignore", invalid="ignore"):
            a_all[region] = (p1 - p0) / np.sqrt(p * (1 - p) * (1 / n_term + 1 / (N - n_term)))
    target_cols = {t: [vocab.index(w) for w in ws] for t, ws in {**TARGETS, **SYMPTOMS}.items()}
    excluded = {c for cs in target_cols.values() for c in cs}
    pool = [c for c in range(len(vocab)) if POOL[0] <= n_term[c] <= POOL[1] and c not in excluded
            and np.isfinite(a_all["anterior_insula"][c])]
    pool = np.array(pool)

    res = {"K": k, "radius_mm": RADIUS, "studies_with_coordinates": int(N), "pool_terms": int(len(pool)),
           "base_rate": {r: round(float(act[r].mean()), 4) for r in act}, "terms": {}}
    for t, cols in target_cols.items():
        res["terms"][t] = {}
        for region in act:
            z, n1, p1, p0 = z_for(cols, region)
            null = a_all[region][pool]
            res["terms"][t][region] = {"z": round(z, 3), "n_studies": n1, "rate_with": round(p1, 4),
                                       "rate_without": round(p0, 4), "percentile_vs_random_terms": round(float((null < z).mean()), 4)}
    zA = {t: res["terms"][t]["anterior_insula"]["z"] for t in target_cols}
    nullA = a_all["anterior_insula"][pool]
    # H1: conjunction (minimum z) of the three functions versus random triples
    obs1 = min(zA[t] for t in TARGETS)
    triples = np.array([nullA[rng.choice(len(pool), 3, replace=False)].min() for _ in range(N_NULL)])
    res["H1_conjunction"] = {"observed_min_z": round(obs1, 3), "percentile": round(float((triples < obs1).mean()), 4),
                             "p_one_sided": round(float(((triples >= obs1).sum() + 1) / (N_NULL + 1)), 4)}
    # H2: anterior minus posterior for each function, versus the same difference for random terms
    diff_null = (a_all["anterior_insula"] - a_all["posterior_insula"])[pool]
    res["H2_anterior_minus_posterior"] = {}
    for t in TARGETS:
        d = res["terms"][t]["anterior_insula"]["z"] - res["terms"][t]["posterior_insula"]["z"]
        res["H2_anterior_minus_posterior"][t] = {"difference_z": round(d, 3), "percentile_vs_random_terms": round(float((diff_null < d).mean()), 4)}
    # H3: mean anterior-insula z of the four symptom terms versus random 4-tuples
    obs3 = float(np.mean([zA[t] for t in SYMPTOMS]))
    quads = np.array([nullA[rng.choice(len(pool), 4, replace=False)].mean() for _ in range(N_NULL)])
    res["H3_symptoms_mean"] = {"observed_mean_z": round(obs3, 3), "percentile": round(float((quads < obs3).mean()), 4),
                               "p_one_sided": round(float(((quads >= obs3).sum() + 1) / (N_NULL + 1)), 4)}
    res["random_term_anterior_insula_z"] = {"median": round(float(np.median(nullA)), 3), "p95": round(float(np.percentile(nullA, 95)), 3)}
    with open(out_path, "w") as fh:
        json.dump(res, fh, indent=2, sort_keys=True)
        fh.write("\n")
    print(json.dumps({k2: v for k2, v in res.items() if k2 != "terms"}, indent=1))


if __name__ == "__main__":
    main()