import json
from collections import defaultdict
from pathlib import Path

import numpy as np
from scipy.stats import spearmanr

RS = Path(__file__).resolve().parents[2] / "data" / "p53" / "rmsd_screen"
COVER = 0.9
N_SPLITS = 400
rng = np.random.default_rng(0)


def load_matrix():
    manifest = json.loads((RS / "receptors.json").read_text())
    recs = {(r["state"], r["frame"]): r for r in manifest}
    per = defaultdict(dict)
    for f in sorted((RS / "scores").glob("*.json")):
        d = json.loads(f.read_text())
        for s in d["scores"]:
            if s.get("affinity") is None:
                continue
            r = recs.get((s.get("state"), s.get("frame")))
            if r:
                per[d["ligand"]][r["frame"]] = float(s["affinity"])
    ligs_all = sorted(per)
    all_frames = sorted({r["frame"] for r in manifest})
    cov = {i: sum(i in per[l] for l in ligs_all) / len(ligs_all) for i in all_frames}
    frames = [i for i in all_frames if cov[i] >= COVER]
    lig_cov = {l: sum(i in per[l] for i in frames) / len(frames) for l in ligs_all}
    ligs = [l for l in ligs_all if lig_cov[l] >= COVER]
    rec_by_frame = {r["frame"]: r for r in manifest}
    vol = np.array([rec_by_frame[i]["site_volume"] for i in frames], dtype=float)
    S = np.full((len(ligs), len(frames)), np.nan)
    for li, l in enumerate(ligs):
        for ri, i in enumerate(frames):
            if i in per[l]:
                S[li, ri] = per[l][i]
    row_mean = np.nanmean(S, axis=1, keepdims=True)
    S = np.where(np.isnan(S), row_mean, S)
    return S, ligs, frames, vol, rec_by_frame, manifest


def volume_trend(Z, vol, nbins=10):
    """Return the pooled volume-decile trend for each receptor."""
    edges = np.quantile(vol, np.linspace(0, 1, nbins + 1))
    edges[-1] += 1e-9
    bin_id = np.clip(np.searchsorted(edges, vol, side="right") - 1, 0, nbins - 1)
    g = np.array(
        [Z[:, bin_id == b].mean() if (bin_id == b).any() else 0.0 for b in range(nbins)]
    )
    return g[bin_id]


def gamma_matrix():
    """Exactly the manuscript's gamma: S -> z-score by ligand -> remove f(V_r) -> remove beta_r."""
    S, ligs, frames, vol, rec_by_frame, manifest = load_matrix()
    Z = (S - S.mean(axis=1, keepdims=True)) / S.std(axis=1, keepdims=True)
    G = volume_trend(Z, vol)
    R = Z - G[None, :]
    gamma = R - R.mean(axis=0)[None, :]
    strat = np.array([rec_by_frame[i]["state"] for i in frames])
    return gamma, np.array(ligs), np.array(frames), strat


def pref_corr(mat):
    """Ligand-preference profile = ligand x ligand correlation over the given receptors."""
    return np.corrcoef(mat)


def split_half_within(gamma, strat, n_splits=N_SPLITS, shuffle_strata=False):
    """Split receptors into halves INSIDE each stratum, then compare ligand-preference
    structure between the halves. Between-stratum contrasts cannot drive this."""
    n_lig = gamma.shape[0]
    iu = np.triu_indices(n_lig, k=1)
    out = []
    for _ in range(n_splits):
        s = strat.copy()
        if shuffle_strata:
            s = rng.permutation(s)
        A, B = [], []
        for lab in np.unique(s):
            idx = np.where(s == lab)[0]
            perm = rng.permutation(idx)
            A.extend(perm[: len(idx) // 2])
            B.extend(perm[len(idx) // 2 :])
        A, B = np.array(A), np.array(B)
        if len(A) < 3 or len(B) < 3:
            continue

        gc = gamma.copy()
        for lab in np.unique(s):
            m = s == lab
            gc[:, m] -= gc[:, m].mean(axis=1, keepdims=True)
        r_, _ = spearmanr(pref_corr(gc[:, A])[iu], pref_corr(gc[:, B])[iu])
        if np.isfinite(r_):
            out.append(r_)
    return np.array(out)