# Reviewer diagnostics on the Allen brain: split-half reliability of subclass excess (odd vs even
# sections), and how much of each subclass's tile its own subclass fills (self-share shrinkage).
import sys, numpy as np, pandas as pd
from scipy.spatial import cKDTree
from scipy.stats import spearmanr
bundle = sys.argv[1]
z = np.load(bundle + "/data/cells_allen.npz", allow_pickle=False)
c = {k: z[k] for k in z.files}
per = pd.read_csv(bundle + "/results/per_subclass.csv", index_col=0)
x = c["x_um"].astype(float); y = c["y_um"].astype(float); kind = c["kind"]; sec = c["sec"].astype(int)
neu = kind == 1
d = np.full(len(kind), np.nan)
for s in np.unique(sec):
    m = sec == s; t = m & (kind == 2); q = m & neu
    if t.any() and q.any():
        d[q] = cKDTree(np.c_[x[t], y[t]]).query(np.c_[x[q], y[q]], k=1)[0]
ok = neu & np.isfinite(d) & (d <= 200)
v = np.log(np.maximum(d, 1.0))
tx = np.floor(x / 100).astype(np.int64); ty = np.floor(y / 100).astype(np.int64)
_, tile = np.unique(np.c_[sec, tx, ty], axis=0, return_inverse=True); tile = tile.ravel()
ssum = np.bincount(tile[ok], v[ok], minlength=tile.max() + 1); n = np.bincount(tile[ok], minlength=tile.max() + 1)
r = np.full(len(d), np.nan); r[ok] = v[ok] - ssum[tile[ok]] / n[tile[ok]]
df = pd.DataFrame({"sub": c["subclass"][ok], "r": r[ok], "sec": sec[ok], "tile": tile[ok]})
full = df.groupby("sub").r.mean()
print("my excess vs bundle's (max abs diff over primary):", float((full.reindex(per.index) - per.excess).abs().max()))
odd = df[df.sec % 2 == 1].groupby("sub").r.agg(["mean", "size"]); even = df[df.sec % 2 == 0].groupby("sub").r.agg(["mean", "size"])
j = per[["E"]].join(odd.add_prefix("o_")).join(even.add_prefix("e_")).dropna()
j = j[(j.o_size >= 100) & (j.e_size >= 100)]
rh = spearmanr(j.o_mean, j.e_mean).statistic
print("split-half units", len(j), "rho(odd, even)", round(rh, 3), "Spearman-Brown full-brain reliability", round(2 * rh / (1 + rh), 3))
print("rho(E, odd)", round(spearmanr(j.E, j.o_mean).statistic, 3), "rho(E, even)", round(spearmanr(j.E, j.e_mean).statistic, 3))
# self-share: fraction of a neuron's tile (included neurons) that is its own subclass
cnt = df.groupby(["tile", "sub"]).size().rename("k").reset_index()
cnt["n"] = n[cnt.tile.values]
cnt["share"] = cnt.k / cnt.n
w = cnt.assign(ws=cnt.k * cnt.share).groupby("sub")[["ws", "k"]].sum()
share = (w.ws / w.k).reindex(per.index)
print("self-share over primary subclasses: median", round(share.median(), 3), "90th pct", round(share.quantile(0.9), 3), "max", round(share.max(), 3), share.idxmax())
print("top self-share:"); print(share.sort_values(ascending=False).head(8).round(3).to_string())
print("rho(E, self-share)", round(spearmanr(per.E, share).statistic, 3))
adj = per.excess / (1 - share)
print("rho(E, excess/(1-share))", round(spearmanr(per.E, adj).statistic, 3))
tiles_one = (n[np.unique(tile[ok])] == 1).sum(); print("tiles with one included neuron:", int(tiles_one), "of", len(np.unique(tile[ok])))
print("median included neurons per tile:", float(np.median(n[np.unique(tile[ok])])))