"""Statistical routines used by the analysis, implemented from first principles.

Every test reported in the paper is computed here rather than taken from a
statistics library. The reason is reproducibility over time: SciPy has changed
the default behaviour of several of these routines between versions (the
continuity correction in the Mann-Whitney test, for one), so a reader running
this archive in five years could otherwise obtain numbers that differ from the
published ones without any way of telling why. The only dependency is NumPy.

Contents
    betai, _betacf   regularized incomplete beta function, by continued fraction
    t_sf2            two-sided Student t tail probability
    norm_sf          upper tail of the standard normal
    pearsonr         Pearson correlation with its t-test
    linregress       ordinary least squares with slope test
    mannwhitneyu     Mann-Whitney U, normal approximation with tie correction

All of these were checked against published critical-value tables; the checks
are part of the verification the authors ran and are described in Sect. 2.8.
"""
import math
import numpy as np


def _betacf(a, b, x, itmax=300, eps=3e-12):
    """Continued-fraction expansion for the incomplete beta function, evaluated
    by the modified Lentz algorithm. The 1e-30 floors guard against a zero
    denominator, which is the standard safeguard for this recurrence."""
    qab, qap, qam = a + b, a + 1.0, a - 1.0
    c, d = 1.0, 1.0 - qab * x / qap
    if abs(d) < 1e-30:
        d = 1e-30
    d = 1.0 / d
    h = d
    for m in range(1, itmax + 1):
        m2 = 2 * m
        aa = m * (b - m) * x / ((qam + m2) * (a + m2))
        d = 1.0 + aa * d
        if abs(d) < 1e-30:
            d = 1e-30
        c = 1.0 + aa / c
        if abs(c) < 1e-30:
            c = 1e-30
        d = 1.0 / d
        h *= d * c
        aa = -(a + m) * (qab + m) * x / ((a + m2) * (qap + m2))
        d = 1.0 + aa * d
        if abs(d) < 1e-30:
            d = 1e-30
        c = 1.0 + aa / c
        if abs(c) < 1e-30:
            c = 1e-30
        d = 1.0 / d
        de = d * c
        h *= de
        if abs(de - 1.0) < eps:
            break
    return h


def betai(a, b, x):
    """Regularized incomplete beta I_x(a,b)."""
    if x <= 0.0:
        return 0.0
    if x >= 1.0:
        return 1.0
    lbeta = (math.lgamma(a + b) - math.lgamma(a) - math.lgamma(b)
             + a * math.log(x) + b * math.log(1.0 - x))
    if x < (a + 1.0) / (a + b + 2.0):
        return math.exp(lbeta) * _betacf(a, b, x) / a
    return 1.0 - math.exp(lbeta) * _betacf(b, a, 1.0 - x) / b


def t_sf2(t, df):
    """Two-sided p-value for Student t."""
    t = abs(float(t))
    if df <= 0:
        return float("nan")
    return betai(df / 2.0, 0.5, df / (df + t * t))


def norm_sf(z):
    """Upper-tail probability of the standard normal, via the complementary
    error function. Used for the normal approximation to Mann-Whitney U."""
    return 0.5 * math.erfc(z / math.sqrt(2.0))


def pearsonr(x, y):
    """Pearson product-moment correlation and its two-sided probability.

    The probability comes from t = r sqrt((n-2)/(1-r^2)) on n-2 degrees of
    freedom. r is clipped to [-1, 1] before use because floating-point error can
    push it a fraction beyond 1 for a near-perfect fit, which would make the
    square root of a negative number; when |r| = 1 exactly the probability is
    returned as zero rather than computed."""
    x = np.asarray(x, float); y = np.asarray(y, float)
    n = len(x)
    xm, ym = x - x.mean(), y - y.mean()
    r = float((xm * ym).sum() / math.sqrt((xm ** 2).sum() * (ym ** 2).sum()))
    r = max(-1.0, min(1.0, r))
    if abs(r) >= 1.0:
        return r, 0.0
    t = r * math.sqrt((n - 2) / (1 - r * r))
    return r, t_sf2(t, n - 2)


def linregress(x, y):
    """Ordinary least squares of y on x.

    Returns (slope, intercept, r, p_slope, se_slope). The probability tests the
    null that the slope is zero, on n-2 degrees of freedom. Used for the
    dose-response of mature-product concentration on litter fraction in
    Sect. 3.1, where the eight resulting probabilities are afterwards adjusted
    for multiple testing by the Holm-Bonferroni procedure."""
    x = np.asarray(x, float); y = np.asarray(y, float)
    n = len(x)
    sxx = ((x - x.mean()) ** 2).sum()
    slope = ((x - x.mean()) * (y - y.mean())).sum() / sxx
    intercept = y.mean() - slope * x.mean()
    resid = y - (intercept + slope * x)
    df = n - 2
    s2 = (resid ** 2).sum() / df
    se = math.sqrt(s2 / sxx)
    r, _ = pearsonr(x, y)
    p = t_sf2(slope / se, df) if se > 0 else 0.0
    return slope, intercept, r, p, se


def mannwhitneyu(a, b):
    """Two-sided Mann-Whitney U test, normal approximation with tie correction.

    Returns (U1, p). Ranks are assigned with ties averaged, and the variance
    carries the standard tie correction sum(t^3 - t)/(N(N-1)). A continuity
    correction of 0.5 is applied to the normal approximation.

    At the sample sizes used here the normal approximation is conservative
    relative to the exact permutation distribution, that is, it returns a larger
    probability than the exact test. The paper states this and quotes the exact
    value alongside it where the contrast is central, so the approximation never
    carries an argument on its own."""
    a = np.asarray(a, float); b = np.asarray(b, float)
    n1, n2 = len(a), len(b)
    allv = np.concatenate([a, b])
    order = allv.argsort()
    ranks = np.empty(len(allv), float)
    sv = allv[order]
    i = 0
    while i < len(sv):
        j = i
        while j + 1 < len(sv) and sv[j + 1] == sv[i]:
            j += 1
        ranks[order[i:j + 1]] = (i + j) / 2.0 + 1.0
        i = j + 1
    R1 = ranks[:n1].sum()
    U1 = R1 - n1 * (n1 + 1) / 2.0
    U2 = n1 * n2 - U1
    U = min(U1, U2)
    mu = n1 * n2 / 2.0
    _, cnt = np.unique(allv, return_counts=True)
    tie = ((cnt ** 3 - cnt).sum())
    N = n1 + n2
    sd = math.sqrt(n1 * n2 / 12.0 * ((N + 1) - tie / (N * (N - 1))))
    z = (U - mu + 0.5) / sd
    return U1, 2 * norm_sf(abs(z))
