"""r_fast.py: vectorised bootstrap statistics for the exploratory analyses (Phase 4).

The frozen functions of q_common.py compute one weighted statistic per call. The exploratory analyses need
the same statistics over thousands of run-bootstrap weight vectors for many feature sets, so these helpers
compute them for a whole weight matrix at once. Each helper is checked against its frozen counterpart on the
first rows of the matrix before it is used (check_against_frozen), so the two can never disagree silently.
"""
import numpy as np
from scipy import sparse

import q_common as QC


def auc_num_den(y, s, W):
    """Numerator and denominator of the weighted Mann-Whitney AUC (midrank ties) for every row of W."""
    y = np.asarray(y); s = np.asarray(s, float); W = np.atleast_2d(np.asarray(W, float))
    u, inv = np.unique(s, return_inverse=True)
    n = len(s)
    M = sparse.csr_matrix((np.ones(n), (np.arange(n), inv)), shape=(n, len(u)))
    pw = np.asarray((M.T @ (W * (y == 1)).T)).T
    nw = np.asarray((M.T @ (W * (y == 0)).T)).T
    below = np.cumsum(nw, axis=1) - nw
    num = (pw * (below + 0.5 * nw)).sum(axis=1)
    den = pw.sum(axis=1) * nw.sum(axis=1)
    return num, den


def ratio(num, den):
    out = np.full(len(num), np.nan)
    ok = den > 0
    out[ok] = num[ok] / den[ok]
    return out


def auc_draws(y, s, W):
    return ratio(*auc_num_den(y, s, W))


def auc_rep_draws(y, S, W):
    """Mean over repeats of each repeat's pooled AUC, for every row of W (q_common.auc_over_repeats)."""
    S = np.atleast_2d(S)
    return np.mean([auc_draws(y, s, W) for s in S], axis=0)


def within_across(y, s, fam, W):
    """Pooled AUC split into label 1 and label 0 pairs within a family and pairs across families."""
    y = np.asarray(y); fam = np.asarray(fam)
    num, den = auc_num_den(y, s, W)
    wn = np.zeros_like(num); wd = np.zeros_like(den)
    for f in np.unique(fam):
        m = fam == f
        a, b = auc_num_den(y[m], np.asarray(s)[m], W[:, m])
        wn += a; wd += b
    return ratio(num, den), ratio(wn, wd), ratio(num - wn, den - wd), ratio(den - wd, den)


def share_draws(ind, W):
    ind = np.asarray(ind, float); W = np.atleast_2d(W)
    den = W.sum(axis=1)
    return ratio(W @ ind, den)


def check_against_frozen(y, s, W, fam=None, rows=3):
    """Assert agreement with q_common.auc_w and q_common.within_family_auc on the first rows of W."""
    got = auc_draws(y, s, W[:rows])
    want = np.array([QC.auc_w(y, s, w) for w in W[:rows]])
    assert np.allclose(got, want, equal_nan=True), (got, want)
    if fam is not None:
        _, wi, _, _ = within_across(y, s, fam, W[:rows])
        want_wi = np.array([QC.within_family_auc(y, s, fam, np.unique(fam), w) for w in W[:rows]])
        assert np.allclose(wi, want_wi, equal_nan=True), (wi, want_wi)
    return True