# (c) 2026 William J. Dwyer. Free and open-source software under the MIT License (see LICENSE).
# If you use or build upon this software, please retain this notice and cite the accompanying
# methodology paper together with the archived software (Zenodo DOI: 10.5281/zenodo.20709963).
"""Recovering parametric inference: can the CV screen license a t-test where
convention uses a non-parametric test?

Compares, under H0 (size) and under a true effect (power):
  raw_t  : Welch t on the natural scale
  log_t  : Welch t on log scale (exact under lognormality)
  perm_t : permutation t on log scale (exact size regardless of distribution)
  MW     : Mann-Whitney U (rank-based, distribution-free)
  adapt  : adaptive PARAMETRIC pipeline -- choose log scale if the pooled sample
           CV > 0.5 (a permutation-invariant feature), else raw scale, then t-test.

Also: Mann-Whitney small-sample resolution floor, and behavior under
misspecification (gamma data that are skewed but NOT lognormal).
"""
import numpy as np, json, warnings
from scipy import stats
warnings.filterwarnings("ignore")
rng = np.random.default_rng(13579)
A = 0.05

def lognormal(cv, n, reps, logshift=0.0):
    s2 = np.log(cv**2 + 1.0); mu_e = -s2/2.0 + logshift  # base median 1
    return rng.lognormal(mu_e, np.sqrt(s2), (reps, n))
def gamma(k, n, reps, logshift=0.0):
    return rng.gamma(k, np.exp(logshift), (reps, n))     # shape k, scale exp(shift); CV=1/sqrt(k)
def posnormal(n, reps):
    return rng.normal(10.0, 2.0, (reps, n))              # positive, CV=0.2

def rates(g1, g2, do_perm=False, B=499):
    raw = stats.ttest_ind(g1, g2, axis=1, equal_var=False).pvalue
    lg  = stats.ttest_ind(np.log(g1), np.log(g2), axis=1, equal_var=False).pvalue
    mw  = np.array([stats.mannwhitneyu(a, b, alternative="two-sided").pvalue for a, b in zip(g1, g2)])
    pooled = np.concatenate([g1, g2], axis=1)
    cvp = pooled.std(axis=1, ddof=1) / pooled.mean(axis=1)
    adapt = np.where(cvp > 0.5, lg, raw)
    out = {"raw_t": float(np.mean(raw < A)), "log_t": float(np.mean(lg < A)),
           "MW": float(np.mean(mw < A)), "adapt": float(np.mean(adapt < A))}
    if do_perm:
        n1 = g1.shape[1]; pr = []
        for a, b in zip(np.log(g1), np.log(g2)):
            obs = abs(a.mean() - b.mean()); pool = np.concatenate([a, b]); c = 0
            for _ in range(B):
                p = rng.permutation(pool)
                if abs(p[:n1].mean() - p[n1:].mean()) >= obs - 1e-12: c += 1
            pr.append((c + 1) / (B + 1))
        out["perm_t"] = float(np.mean(np.array(pr) < A))
    return out

results = {}

# ---------- 1) SIZE under H0 (no difference) ----------
ns = [5, 10, 20, 30]; REPS = 3000
configs = {
    "normal_CV0.2": lambda n: (posnormal(n, REPS), posnormal(n, REPS)),
    "lognormal_CV1.0": lambda n: (lognormal(1.0, n, REPS), lognormal(1.0, n, REPS)),
    "lognormal_CV2.0": lambda n: (lognormal(2.0, n, REPS), lognormal(2.0, n, REPS)),
    "gamma_k2_CV0.71": lambda n: (gamma(2.0, n, REPS), gamma(2.0, n, REPS)),
    "gamma_k0.5_CV1.41": lambda n: (gamma(0.5, n, REPS), gamma(0.5, n, REPS)),
}
size = {}
for name, gen in configs.items():
    size[name] = {}
    for n in ns:
        g1, g2 = gen(n)
        size[name][n] = rates(g1, g2)
results["size"] = size

# ---------- 2) POWER under a true effect (median ratio 2 => log shift ln 2) ----------
shift = np.log(2.0); REPS_P = 3000
power = {}
for cv in [1.0, 2.0]:
    power[cv] = {}
    for n in [5, 10, 20, 30, 50]:
        g1 = lognormal(cv, n, REPS_P, 0.0)
        g2 = lognormal(cv, n, REPS_P, shift)
        power[cv][n] = rates(g1, g2)
results["power_logshift_ln2"] = power

# ---------- 3) Mann-Whitney small-sample resolution floor ----------
# Minimum attainable two-sided p when the two groups are completely separated.
floor = {}
for n in range(2, 8):
    a = np.arange(1, n + 1); b = np.arange(n + 1, 2 * n + 1)
    floor[n] = float(stats.mannwhitneyu(a, b, alternative="two-sided").pvalue)
results["mw_min_p_by_n"] = floor

# ---------- 4) Misspecification: gamma (skewed, NOT lognormal), with permutation backstop ----------
mis = {}
for n in [10, 20]:
    g1 = gamma(0.5, n, 1500); g2 = gamma(0.5, n, 1500)   # H0, heavy right skew, not lognormal
    mis[n] = rates(g1, g2, do_perm=True, B=499)
results["misspecification_gamma_k0.5_H0"] = mis

json.dump(results, open("recover.json", "w"), indent=2)

# ---- console ----
m = ["raw_t", "log_t", "MW", "adapt"]
print("SIZE under H0 (reject rate, nominal 0.05):")
for name, d in size.items():
    print(f"\n{name}")
    print("  n   " + "  ".join(f"{t:>6}" for t in m))
    for n in ns:
        print(f"  {n:>3} " + "  ".join(f"{d[n][t]:>6.3f}" for t in m))
print("\nPOWER (median ratio 2):")
for cv in [1.0, 2.0]:
    print(f"\nCV={cv}")
    print("  n   " + "  ".join(f"{t:>6}" for t in ["raw_t","log_t","MW"]))
    for n in [5,10,20,30,50]:
        print(f"  {n:>3} " + "  ".join(f"{power[cv][n][t]:>6.3f}" for t in ["raw_t","log_t","MW"]))
print("\nMANN-WHITNEY minimum two-sided p by n/group:")
for n, p in floor.items(): print(f"  n={n}: min p = {p:.3f}")
print("\nMISSPECIFICATION (gamma k=0.5, H0): size")
for n, d in mis.items():
    print(f"  n={n}: " + ", ".join(f"{t}={d[t]:.3f}" for t in ["raw_t","log_t","perm_t","MW"]))
print("\ndone")
