"""False-positive rates of outlier-exclusion strategies when there is no true difference.

Two groups of n animals, Welch's t-test (the site's default), two-sided alpha = 0.05.
Strategies:
  prespecified   - analyse every value (no exclusion)
  drop_if_ns     - if p > 0.05, drop the single most extreme value (largest |z| within its group) and test
                   again; a false positive if either test gives p < 0.05
  z2_if_ns       - if p > 0.05, drop every value with |z| > 2 within its group, then test again
  any_one        - try dropping each single value in turn; a false positive if any version gives p < 0.05
  mad_always     - before any p-value is seen, drop values more than 3 scaled MADs (MAD x 1.4826) from their
                   group median (the same rule for every dataset), then test once
  switch_if_ns   - no deletion; if Welch p > 0.05, try an exact Mann-Whitney test, then Yuen's 20% trimmed test,
                   then Welch on log values (log only for positive data), keeping the first significant result
Distributions: normal, and lognormal (sigma = 0.8, right-skewed like many concentrations).
The Welch test is computed in vectorised NumPy; check_against_scipy() confirms it matches
scipy.stats.ttest_ind(equal_var=False) on 2,000 random datasets before the simulation runs.
Python 3.13, NumPy 2.5, SciPy 1.18.1. Seed 20261110 (fixed in advance).
Run from this folder: python simulation.py
"""
import json

import numpy as np
from scipy import stats

RNG = np.random.default_rng(20261110)
RUNS = 100_000
ALPHA = 0.05


def welch_p(a, b, ma, mb):
    """Two-sided Welch p for each row; ma/mb are 0/1 masks of the values kept."""
    na, nb = ma.sum(1), mb.sum(1)
    mean_a = (a * ma).sum(1) / na
    mean_b = (b * mb).sum(1) / nb
    va = (((a - mean_a[:, None]) ** 2) * ma).sum(1) / (na - 1)
    vb = (((b - mean_b[:, None]) ** 2) * mb).sum(1) / (nb - 1)
    se2 = va / na + vb / nb
    t = (mean_a - mean_b) / np.sqrt(se2)
    df = se2**2 / ((va / na) ** 2 / (na - 1) + (vb / nb) ** 2 / (nb - 1))
    return 2 * stats.t.sf(np.abs(t), df)


def zscores(x, m):
    n = m.sum(1)
    mean = (x * m).sum(1) / n
    sd = np.sqrt((((x - mean[:, None]) ** 2) * m).sum(1) / (n - 1))
    return (x - mean[:, None]) / sd[:, None]


def mad_mask(x, k=3.0):
    med = np.median(x, axis=1, keepdims=True)
    mad = 1.4826 * np.median(np.abs(x - med), axis=1, keepdims=True)
    with np.errstate(divide="ignore", invalid="ignore"):
        keep = np.where(mad > 0, np.abs(x - med) / mad <= k, True)
    return keep.astype(float)


def check_against_scipy():
    a, b = RNG.normal(size=(2000, 6)), RNG.lognormal(sigma=0.8, size=(2000, 7))
    mine = welch_p(a, b, np.ones_like(a), np.ones_like(b))
    ref = stats.ttest_ind(a, b, axis=1, equal_var=False).pvalue
    return float(np.max(np.abs(mine - ref)))


def mw_null_sf(n1, n2):
    """Exact null distribution of U (no ties): P(U >= u) for u = 0..n1*n2."""
    counts = np.zeros((n1 + 1, n2 + 1, n1 * n2 + 1))
    for i in range(n1 + 1):
        for j in range(n2 + 1):
            if i == 0 or j == 0:
                counts[i, j, 0] = 1
            else:
                counts[i, j] = np.concatenate([np.zeros(j), counts[i - 1, j, : n1 * n2 + 1 - j]]) + counts[i, j - 1]
    pmf = counts[n1, n2] / counts[n1, n2].sum()
    return pmf[::-1].cumsum()[::-1]


def mw_p(a, b):
    """Two-sided exact Mann-Whitney p for each row (continuous data, no ties), as SciPy method='exact'."""
    n1, n2 = a.shape[1], b.shape[1]
    ranks = stats.rankdata(np.concatenate([a, b], axis=1), axis=1)
    u1 = ranks[:, :n1].sum(1) - n1 * (n1 + 1) / 2
    big = np.maximum(u1, n1 * n2 - u1).round().astype(int)
    return np.minimum(1.0, 2 * mw_null_sf(n1, n2)[big])


def yuen_p(a, b, tr=0.2):
    def part(x):
        n = x.shape[1]
        g = int(np.floor(tr * n))
        h = n - 2 * g
        s = np.sort(x, axis=1)
        w = s.copy()
        if g > 0:
            w[:, :g] = s[:, [g]]
            w[:, n - g:] = s[:, [n - g - 1]]
        tm = s[:, g:n - g].mean(1)
        d = (n - 1) * w.var(1, ddof=1) / (h * (h - 1))
        return tm, d, h
    ta, da, ha = part(a)
    tb, db, hb = part(b)
    t = (ta - tb) / np.sqrt(da + db)
    df = (da + db) ** 2 / (da**2 / (ha - 1) + db**2 / (hb - 1))
    return 2 * stats.t.sf(np.abs(t), df)


def check_rank_and_yuen():
    rng = np.random.default_rng(1)
    a, b = rng.normal(size=(300, 6)), rng.lognormal(sigma=0.8, size=(300, 6))
    ref_mw = np.array([stats.mannwhitneyu(x, y, method="exact").pvalue for x, y in zip(a, b)])
    ref_yu = stats.ttest_ind(a, b, axis=1, equal_var=False, trim=0.2).pvalue
    return float(np.max(np.abs(mw_p(a, b) - ref_mw))), float(np.max(np.abs(yuen_p(a, b) - ref_yu)))


def run(n, dist):
    if dist == "normal":
        a, b = RNG.normal(size=(RUNS, n)), RNG.normal(size=(RUNS, n))
    else:
        a, b = RNG.lognormal(sigma=0.8, size=(RUNS, n)), RNG.lognormal(sigma=0.8, size=(RUNS, n))
    ones = np.ones((RUNS, n))
    sig0 = welch_p(a, b, ones, ones) < ALPHA

    za, zb = np.abs(zscores(a, ones)), np.abs(zscores(b, ones))
    drop_a = za.max(1) >= zb.max(1)
    ma1, mb1 = ones.copy(), ones.copy()
    rows = np.arange(RUNS)
    ma1[rows[drop_a], za.argmax(1)[drop_a]] = 0
    mb1[rows[~drop_a], zb.argmax(1)[~drop_a]] = 0
    sig_drop = sig0 | (welch_p(a, b, ma1, mb1) < ALPHA)

    ma2, mb2 = (za <= 2).astype(float), (zb <= 2).astype(float)
    dropped = (ma2.sum(1) < n) | (mb2.sum(1) < n)
    with np.errstate(divide="ignore", invalid="ignore"):
        p2 = welch_p(a, b, ma2, mb2)
    sig_z2 = sig0 | (dropped & (p2 < ALPHA))

    sig_any = sig0.copy()
    for i in range(n):
        m = ones.copy()
        m[:, i] = 0
        sig_any |= welch_p(a, b, m, ones) < ALPHA
        sig_any |= welch_p(a, b, ones, m) < ALPHA

    sig_mad = welch_p(a, b, mad_mask(a), mad_mask(b)) < ALPHA
    sig_switch = sig0 | (mw_p(a, b) < ALPHA) | (yuen_p(a, b) < ALPHA)
    if dist == "lognormal":
        sig_switch |= welch_p(np.log(a), np.log(b), ones, ones) < ALPHA
    share_mad_dropped = float(((mad_mask(a).sum(1) < n) | (mad_mask(b).sum(1) < n)).mean())
    return {
        "prespecified": round(float(sig0.mean()), 4),
        "drop_if_ns": round(float(sig_drop.mean()), 4),
        "z2_if_ns": round(float(sig_z2.mean()), 4),
        "any_one": round(float(sig_any.mean()), 4),
        "mad_always": round(float(sig_mad.mean()), 4),
        "switch_if_ns": round(float(sig_switch.mean()), 4),
        "share_datasets_mad_dropped_any": round(share_mad_dropped, 4),
    }


out = {"runs": RUNS, "seed": 20261110, "alpha": ALPHA, "test": "Welch t-test, two-sided",
       "max_abs_difference_vs_scipy": check_against_scipy(), "rows": []}
for dist, sizes in (("normal", (6, 10, 30)), ("lognormal", (6, 10))):
    for n in sizes:
        out["rows"].append({"distribution": dist, "n_per_group": n, **run(n, dist)})
out["max_abs_difference_vs_scipy_mannwhitney_yuen"] = check_rank_and_yuen()
# Largest possible |z| in a sample of n (Shiffler 1988): (n - 1) / sqrt(n)
out["max_abs_z"] = {str(n): round(float((n - 1) / np.sqrt(n)), 3) for n in (3, 4, 5, 6, 8, 10, 12)}
print(json.dumps(out, indent=2))
with open("simulation_output.json", "w") as f:
    json.dump(out, f, indent=2)