#!/usr/bin/env python
import argparse
import json
import sys
from pathlib import Path

import numpy as np
from scipy.cluster.hierarchy import linkage, leaves_list
from scipy.spatial.distance import squareform
from scipy.stats import spearmanr

ROOT = Path(__file__).resolve().parents[2]
CACHE = ROOT / "data" / "p53" / "fig_data"
sys.path.insert(0, str(ROOT))

SEED = 0


def split_half_all(gamma, n_splits, rng, permute=False):
    """Split-half over ALL receptors: does the ligand-preference structure reproduce?"""
    n_lig, n_rec = gamma.shape
    iu = np.triu_indices(n_lig, k=1)
    out = []
    for _ in range(n_splits):
        m = gamma.copy()
        if permute:
            for li in range(n_lig):
                m[li] = m[li, rng.permutation(n_rec)]
        p = rng.permutation(n_rec)
        A, B = p[: n_rec // 2], p[n_rec // 2 :]
        r_, _ = spearmanr(np.corrcoef(m[:, A])[iu], np.corrcoef(m[:, B])[iu])
        if np.isfinite(r_):
            out.append(r_)
    return np.array(out)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--splits", type=int, default=400)
    a = ap.parse_args()

    from scoobdoob.validation import preference as WS

    CACHE.mkdir(parents=True, exist_ok=True)
    gamma, ligs, frames, strat = WS.gamma_matrix()
    n_lig, n_rec = gamma.shape
    labs = np.unique(strat)

    C = np.corrcoef(gamma)
    D = np.clip(1.0 - C, 0, 2)
    np.fill_diagonal(D, 0.0)
    lig_order = leaves_list(linkage(squareform(D, checks=False), method="average"))
    rec_order = np.concatenate([np.where(strat == l)[0] for l in labs])

    tot = gamma.var()
    between = np.zeros_like(gamma)
    for l in labs:
        m_ = strat == l
        between[:, m_] = gamma[:, m_].mean(axis=1, keepdims=True)
    frac_between = float(between.var() / tot)
    frac_within = float((gamma - between).var() / tot)

    obs_all = split_half_all(gamma, a.splits, np.random.default_rng(SEED))
    nul_all = split_half_all(
        gamma, a.splits, np.random.default_rng(SEED + 1), permute=True
    )

    WS.rng = np.random.default_rng(SEED + 2)
    obs_w = WS.split_half_within(gamma, strat, n_splits=a.splits)
    rng = np.random.default_rng(SEED + 3)
    gperm = gamma.copy()
    for l in labs:
        m_ = np.where(strat == l)[0]
        for li in range(n_lig):
            gperm[li, m_] = gperm[li, rng.permutation(m_)]
    WS.rng = np.random.default_rng(SEED + 4)
    nul_w = WS.split_half_within(gperm, strat, n_splits=a.splits)

    p_within = float((nul_w >= obs_w.mean()).mean())

    np.savez_compressed(
        CACHE / "preference.npz",
        gamma=gamma,
        strat=strat,
        ligands=np.array(ligs),
        frames=np.array(frames),
        lig_order=lig_order,
        rec_order=rec_order,
        obs_all=obs_all,
        nul_all=nul_all,
        obs_within=obs_w,
        nul_within=nul_w,
        frac_within=np.array(frac_within),
        frac_between=np.array(frac_between),
    )
    stats = {
        "n_ligands": int(n_lig),
        "n_receptors": int(n_rec),
        "n_strata": int(len(labs)),
        "stratum_sizes": [int((strat == l).sum()) for l in labs],
        "n_splits": int(a.splits),
        "seed": SEED,
        "gamma_var_total": float(tot),
        "frac_within": frac_within,
        "frac_between": frac_between,
        "splithalf_all_mean": float(obs_all.mean()),
        "splithalf_all_sd": float(obs_all.std()),
        "splithalf_all_null_mean": float(nul_all.mean()),
        "splithalf_within_mean": float(obs_w.mean()),
        "splithalf_within_sd": float(obs_w.std()),
        "splithalf_within_null_mean": float(nul_w.mean()),
        "p_null_ge_observed_within": p_within,
    }
    (CACHE / "preference.json").write_text(json.dumps(stats, indent=1, default=float))

    print(
        f"  {n_lig} ligands x {n_rec} receptors, {len(labs)} strata "
        f"{stats['stratum_sizes']}"
    )
    print(
        f"  var(gamma): within {100 * frac_within:.1f}%  between {100 * frac_between:.1f}%"
    )
    print(
        f"  split-half all receptors : {obs_all.mean():+.4f} +/- {obs_all.std():.4f} "
        f"(null {nul_all.mean():+.4f})"
    )
    print(
        f"  split-half within stratum: {obs_w.mean():+.4f} +/- {obs_w.std():.4f} "
        f"(null {nul_w.mean():+.4f}, p={p_within:.4f})"
    )
    print(f"-> {CACHE / 'preference.npz'}\n-> {CACHE / 'preference.json'}")


if __name__ == "__main__":
    main()