#!/usr/bin/env python3
"""Preregistered analysis: do energy-hungry neuron types sit closer to capillaries?
usage: analyze.py BUNDLE_DIR [--energy-only | --pilot BRAIN SECTION_INDEX]"""
import sys, json, hashlib, pathlib, time
import numpy as np, pandas as pd
from scipy.spatial import cKDTree
from scipy.stats import rankdata

ROOT = pathlib.Path(sys.argv[1]); D = ROOT / "data"; R = ROOT / "results"
P = json.load(open(D / "parameters.json"))
plan_p = ROOT / "plan" / "parameters.json"
if plan_p.exists():
    assert plan_p.read_bytes() == (D / "parameters.json").read_bytes(), "data/parameters.json differs from plan/parameters.json"
BRAINS = [P["primary_brain"]] + P["replication_brains"]

def rnd(x, k=8):
    if isinstance(x, (float, np.floating)):
        return float(f"{float(x):.{k}g}") if np.isfinite(x) else None
    if isinstance(x, (np.integer,)): return int(x)
    if isinstance(x, dict): return {a: rnd(b, k) for a, b in x.items()}
    if isinstance(x, (list, tuple)): return [rnd(b, k) for b in x]
    if isinstance(x, (np.bool_,)): return bool(x)
    return x

def zs(v):
    v = np.asarray(v, float); return (v - v.mean()) / v.std()

# ---------------- energy score (locked) ----------------
def energy(extra_subclasses=()):
    sub = pd.read_csv(D / "subclasses.csv", index_col=0)
    prim = list(sub.index[sub.primary]) + [s for s in extra_subclasses if s not in sub.index[sub.primary]]
    x = pd.read_csv(D / "subclass_expression.csv", index_col=0)
    eg = pd.read_csv(D / "energy_genes.csv"); eg = eg[eg.in_10x]
    X = x.loc[prim, list(eg.symbol)]
    base = x.loc[list(sub.index[sub.primary]), list(eg.symbol)]
    expressed = base.max() >= P["gene_min_max_subclass_mean_log2"]
    ox = [g for g, c in zip(eg.symbol, eg.component) if c == "oxphos" and expressed[g]]
    pu = [g for g, c in zip(eg.symbol, eg.component) if c == "pump" and expressed[g]]
    mu, sd = base.mean(), base.std(ddof=0)
    Z = (X - mu) / sd
    oxs, pus = Z[ox].mean(axis=1), Z[pu].mean(axis=1)
    b_ox, b_pu = oxs.loc[base.index], pus.loc[base.index]
    E = ((oxs - b_ox.mean()) / b_ox.std(ddof=0) + (pus - b_pu.mean()) / b_pu.std(ddof=0)) / 2
    out = pd.DataFrame({"E": E, "oxphos": oxs, "pump": pus, "class": sub["class"].reindex(prim)})
    return out, ox, pu, sub

# ---------------- distances ----------------
def load(brain):
    z = np.load(D / f"cells_{brain}.npz", allow_pickle=False)
    return {k: z[k] for k in z.files}

def nn_dist(c, target_kind, neuron_mask):
    d = np.full(len(c["kind"]), np.nan)
    x = c["x_um"].astype(np.float64); y = c["y_um"].astype(np.float64)
    for s in np.unique(c["sec"]):
        m = c["sec"] == s
        t = m & (c["kind"] == target_kind); q = m & neuron_mask
        if t.sum() == 0 or q.sum() == 0: continue
        tree = cKDTree(np.c_[x[t], y[t]])
        dd, _ = tree.query(np.c_[x[q], y[q]], k=1)
        d[q] = dd
    return d

def residuals(c, d, pool, tile, transform="log"):
    """exact expectation of the within-(section x tile) label-permutation null"""
    ok = pool & np.isfinite(d) & (d <= P["distance_cap_um"])
    v = np.log(np.maximum(d, P["distance_floor_um"])) if transform == "log" else np.maximum(d, P["distance_floor_um"])
    tx = np.floor(c["x_um"].astype(np.float64) / tile).astype(np.int64); ty = np.floor(c["y_um"].astype(np.float64) / tile).astype(np.int64)
    key = np.c_[c["sec"].astype(np.int64), tx, ty]
    _, inv = np.unique(key, axis=0, return_inverse=True); inv = inv.ravel()
    s = np.bincount(inv[ok], weights=v[ok], minlength=inv.max() + 1)
    n = np.bincount(inv[ok], minlength=inv.max() + 1)
    r = np.full(len(d), np.nan)
    r[ok] = v[ok] - s[inv[ok]] / n[inv[ok]]
    return r, ok

def excess_by(c, r, ok, groups):
    df = pd.DataFrame({g: c[g] for g in groups}); df["r"] = r
    df = df[ok]
    g = df.groupby(groups, sort=True)["r"]
    return pd.DataFrame({"excess": g.mean(), "n": g.size()})

# ---------------- statistics ----------------
def rk(v): return rankdata(v, method="average")
def pear(a, b):
    a = a - a.mean(); b = b - b.mean(); return float((a * b).sum() / np.sqrt((a * a).sum() * (b * b).sum()))
def spearman(a, b): return pear(rk(a), rk(b))

def perm_p(a, b, rng, n_perm, strata=None):
    ra, rb = rk(a), rk(b)
    if strata is not None:
        ra, rb = strat_ranks(a, strata), strat_ranks(b, strata)
    obs = pear(ra, rb); cnt = 0
    for _ in range(n_perm):
        if strata is None:
            pa = ra[rng.permutation(len(ra))]
        else:
            pa = ra.copy()
            for s in np.unique(strata):
                ix = np.where(strata == s)[0]; pa[ix] = ra[ix[rng.permutation(len(ix))]]
        cnt += abs(pear(pa, rb)) >= abs(obs) - 1e-12
    return obs, (cnt + 1) / (n_perm + 1)

def boot_ci(a, b, rng, n_boot, strata=None):
    vals = np.empty(n_boot); n = len(a)
    for i in range(n_boot):
        ix = rng.integers(0, n, n)
        if strata is None: vals[i] = spearman(a[ix], b[ix])
        else:
            ra, rb = strat_ranks(a[ix], strata[ix]), strat_ranks(b[ix], strata[ix])
            vals[i] = pear(ra, rb) if ra.std() > 0 and rb.std() > 0 else np.nan
    return float(np.nanpercentile(vals, 2.5)), float(np.nanpercentile(vals, 97.5))

def strat_ranks(v, strata):
    out = np.zeros(len(v))
    for s in np.unique(strata):
        ix = strata == s; r = rk(v[ix]); out[ix] = r - r.mean()
    return out

def corr_block(a, b, seed, strata=None):
    rng = np.random.default_rng(seed)
    rho, p = perm_p(a, b, rng, P["n_perm"], strata)
    lo, hi = boot_ci(a, b, np.random.default_rng(seed + 1), P["n_boot"], strata)
    return {"rho": rho, "p_perm_two_sided": p, "ci95_low": lo, "ci95_high": hi, "n_units": int(len(a))}

def partial_spearman(a, b, c, seed):
    ra, rb, rc = rk(a), rk(b), rk(c)
    X = np.c_[np.ones(len(rc)), rc]
    ea = ra - X @ np.linalg.lstsq(X, ra, rcond=None)[0]; eb = rb - X @ np.linalg.lstsq(X, rb, rcond=None)[0]
    obs = pear(ea, eb); rng = np.random.default_rng(seed); cnt = 0
    for _ in range(P["n_perm"]):
        cnt += abs(pear(ea[rng.permutation(len(ea))], eb)) >= abs(obs) - 1e-12
    return {"partial_rho": obs, "p_perm_two_sided": (cnt + 1) / (P["n_perm"] + 1), "n_units": int(len(a))}

# ---------------- modes ----------------
def energy_only():
    E, ox, pu, sub = energy()
    print(json.dumps(rnd({"n_primary_subclasses": len(E), "n_oxphos_genes_used": len(ox), "pump_genes_used": pu,
                          "oxphos_genes_used": ox, "corr_oxphos_pump_spearman": spearman(E.oxphos.values, E.pump.values),
                          "E_by_class_mean": E.groupby("class").E.mean().round(3).to_dict()}), indent=1))

def pilot(brain, si):
    t0 = time.time(); c = load(brain); keep = c["sec"] == int(si)
    c = {k: (v[keep] if k != "section_names" else v) for k, v in c.items()}
    neu = c["kind"] == 1
    d = nn_dist(c, 2, neu); dp = nn_dist(c, 3, neu)
    r, ok = residuals(c, d, neu, P["tile_um"])
    q = lambda v: [float(np.nanpercentile(v[neu], p)) for p in (10, 50, 90)]
    print(json.dumps(rnd({"pilot": f"{brain} section {c['section_names'][int(si)]}", "neurons": int(neu.sum()),
        "endothelial": int((c['kind'] == 2).sum()), "pericytes": int((c['kind'] == 3).sum()),
        "endo_dist_p10_p50_p90_um": q(d), "peri_dist_p10_p50_p90_um": q(dp),
        "beyond_cap": int((d[neu] > P["distance_cap_um"]).sum()), "residual_mean_abs": float(np.nanmean(np.abs(r))),
        "seconds": time.time() - t0}), indent=1))

def main():
    t0 = time.time(); R.mkdir(exist_ok=True)
    E, ox, pu, sub = energy()
    prim = list(E.index); res = {}; sec = {}; qc = {}
    per = {}
    for b in BRAINS:
        c = load(b); neu = c["kind"] == 1; neu4 = (c["kind"] == 1) | (c["kind"] == 4)
        d = nn_dist(c, 2, neu4); dp = nn_dist(c, 3, neu4)
        out = {}
        r, ok = residuals(c, d, neu, P["tile_um"]); out["primary"] = (r, ok)
        qc[b] = {"neurons": int(neu.sum()), "neurons_analysed": int(ok.sum()), "beyond_cap_or_no_vessel": int((neu & ~ok).sum()),
                 "endothelial": int((c["kind"] == 2).sum()), "pericytes": int((c["kind"] == 3).sum()), "sections": int(len(c["section_names"])),
                 "endo_dist_median_um": float(np.nanmedian(d[neu])), "peri_dist_median_um": float(np.nanmedian(dp[neu]))}
        ex = excess_by(c, r, ok, ["subclass"]); per[(b, "primary")] = ex
        if b == P["primary_brain"] or True:
            per[(b, "pericyte")] = excess_by(c, *residuals(c, dp, neu, P["tile_um"]), ["subclass"])
        if b == P["primary_brain"]:
            per[(b, "linear")] = excess_by(c, *residuals(c, d, neu, P["tile_um"], "linear"), ["subclass"])
            for t in P["tile_um_sensitivity"]:
                per[(b, f"tile{int(t)}")] = excess_by(c, *residuals(c, d, neu, t), ["subclass"])
            per[(b, "division")] = excess_by(c, r, ok, ["subclass", "division"])
            per[(b, "with_imn")] = excess_by(c, *residuals(c, d, neu4, P["tile_um"]), ["subclass"])
        if "log10_counts" in c:
            df = pd.DataFrame({"subclass": c["subclass"][neu], "lc": c["log10_counts"][neu]})
            per[(b, "size")] = df.groupby("subclass").lc.agg(["sum", "size"])
        print(b, "done", round(time.time() - t0), "s", flush=True)

    A = per[(P["primary_brain"], "primary")]
    tab = E.join(A, how="left")
    assert tab.n.min() >= 1
    tab["n_analysed"] = tab.n.astype(int)
    # primary
    pr = corr_block(tab.E.values, tab.excess.values, P["seed"])
    reps = {}
    for b in P["replication_brains"]:
        x = E.join(per[(b, "primary")], how="inner"); x = x[x.n >= P["min_neurons_replication"]]
        reps[b] = {"rho": spearman(x.E.values, x.excess.values), "n_units": int(len(x))}
        tab[f"excess_{b}"] = per[(b, "primary")].excess.reindex(tab.index); tab[f"n_{b}"] = per[(b, "primary")].n.reindex(tab.index)
    n_neg = sum(v["rho"] < 0 for v in reps.values())
    support = pr["rho"] <= P["support_rho_max"] and pr["p_perm_two_sided"] < P["alpha"] and n_neg >= P["replication_min_negative"]
    refute = pr["ci95_low"] > P["refute_ci_low_min"]
    decision = "support" if support else ("refute" if refute else "inconclusive")
    res = {**pr, "replication": reps, "replication_n_negative": n_neg, "decision": decision,
           "supports_energy_proximity": decision == "support", "n_primary_subclasses": len(tab),
           "neurons_analysed_primary_subclasses": int(tab.n_analysed.sum())}
    # secondaries
    s = {}
    pe = per[(P["primary_brain"], "pericyte")].reindex(tab.index); tab["excess_pericyte"] = pe.excess
    s["S1_pericyte"] = corr_block(tab.E.values, pe.excess.values, P["seed"] + 10)
    s["S2_class_stratified"] = corr_block(tab.E.values, tab.excess.values, P["seed"] + 20, strata=tab["class"].values)
    s["S2_class_stratified"]["n_strata_used"] = int((tab["class"].value_counts() >= P["min_units_per_stratum"]).sum())
    dv = per[(P["primary_brain"], "division")].reset_index()
    dv = dv[(dv.n >= P["min_neurons_subclass_division"]) & dv.subclass.isin(tab.index)]
    vc = dv.division.value_counts(); dv = dv[dv.division.map(vc) >= P["min_units_per_stratum"]].sort_values(["division", "subclass"])
    s["S3_within_division"] = corr_block(tab.E.reindex(dv.subclass).values, dv.excess.values, P["seed"] + 30, strata=dv.division.values)
    s["S3_within_division"]["n_divisions"] = int(dv.division.nunique())
    s["S3_within_division"]["per_division_rho"] = {k: spearman(tab.E.reindex(g.subclass).values, g.excess.values) for k, g in dv.groupby("division")}
    sz = sum(per[(b, "size")]["sum"].reindex(tab.index).fillna(0) for b in P["replication_brains"])
    sn = sum(per[(b, "size")]["size"].reindex(tab.index).fillna(0) for b in P["replication_brains"])
    tab["size_log10_counts_zhuang"] = (sz / sn).where(sn >= 50)
    m = tab.size_log10_counts_zhuang.notna()
    s["S4_size_partial"] = partial_spearman(tab.E.values[m], tab.excess.values[m], tab.size_log10_counts_zhuang.values[m], P["seed"] + 40)
    s["S4_size_partial"]["rho_E_size"] = spearman(tab.E.values[m], tab.size_log10_counts_zhuang.values[m])
    s["S4_size_partial"]["rho_excess_size"] = spearman(tab.excess.values[m], tab.size_log10_counts_zhuang.values[m])
    s["S5_oxphos_only"] = corr_block(tab.oxphos.values, tab.excess.values, P["seed"] + 50)
    s["S5_pump_only"] = corr_block(tab.pump.values, tab.excess.values, P["seed"] + 51)
    li = per[(P["primary_brain"], "linear")].reindex(tab.index); tab["excess_linear_um"] = li.excess
    s["S6_linear_distance"] = corr_block(tab.E.values, li.excess.values, P["seed"] + 60)
    for t in P["tile_um_sensitivity"]:
        x = per[(P["primary_brain"], f"tile{int(t)}")].reindex(tab.index)
        s[f"S7_tile_{int(t)}um"] = corr_block(tab.E.values, x.excess.values, P["seed"] + 70 + int(t))
    pool = pd.concat([per[(b, "primary")].assign(w=lambda x: x.excess * x.n) for b in P["replication_brains"]]).groupby(level=0)[["w", "n"]].sum()
    pool = pool[pool.n >= P["min_neurons_replication"]]; pool["excess"] = pool.w / pool.n
    x = E.join(pool[["excess"]], how="inner")
    s["S8_zhuang_pooled"] = corr_block(x.E.values, x.excess.values, P["seed"] + 80)
    s["S8_zhuang_pericyte_per_brain"] = {b: spearman(*E.join(per[(b, "pericyte")], how="inner").query("n >= @P['min_neurons_replication']")[["E", "excess"]].values.T) for b in P["replication_brains"]}
    # IMN sensitivity
    sub = pd.read_csv(D / "subclasses.csv", index_col=0)
    imn = [i for i in sub.index if P["immature_subclass_token"] in i.split() and sub.loc[i, "allen_neurons"] >= 0]
    E2, _, _, _ = energy(extra_subclasses=imn)
    wi = per[(P["primary_brain"], "with_imn")]
    x = E2.join(wi, how="inner"); x = x[x.n >= P["min_neurons_primary"]]
    s["S9_include_immature_neurons"] = corr_block(x.E.values, x.excess.values, P["seed"] + 90)
    # exploratory
    tab["excess_rank_pct"] = rk(tab.excess.values) / len(tab) * 100
    s["X1_hypoxia_vulnerable"] = {k: {"excess": float(tab.loc[k, "excess"]), "pct_rank_excess": float(tab.loc[k, "excess_rank_pct"]),
                                      "E": float(tab.loc[k, "E"])} for k in P["exploratory_subclasses"] if k in tab.index}
    s["X2_class_mean_excess"] = tab.groupby("class").excess.mean().to_dict()
    json.dump(rnd(res), open(R / "R1.json", "w"), indent=1, sort_keys=True)
    json.dump(rnd(s), open(R / "secondary.json", "w"), indent=1, sort_keys=True)
    json.dump(rnd({**qc, "oxphos_genes_used": ox, "pump_genes_used": pu}), open(R / "qc.json", "w"), indent=1, sort_keys=True)
    tab.sort_index().to_csv(R / "per_subclass.csv", float_format="%.8g")
    print(json.dumps(rnd(res), indent=1)); print("total", round(time.time() - t0), "s")

if __name__ == "__main__":
    if "--energy-only" in sys.argv: energy_only()
    elif "--pilot" in sys.argv: i = sys.argv.index("--pilot"); pilot(sys.argv[i + 1], sys.argv[i + 2])
    else: main()