# (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).
"""
Simulation study for: "Finding a cutoff to detect non-normal distributions
from summary statistics" (WJ Dwyer).

Underlying model (paper eq 2): x = exp(mu_e + sigma_e * N(0,1))  -> LOGNORMAL.
Natural-domain mean/SD relate to log-domain params by (paper eqs 8,10,13):
    CV^2 = sigma_n^2 / mu_n^2 = exp(sigma_e^2) - 1
    sigma_e^2 = ln(CV^2 + 1)
    mu_e      = ln(mu_n) - sigma_e^2 / 2     (== 0.5*ln{mu_n^2/(CV^2+1)})

Key identity used throughout: the paper's eq-17 flag
    sigma_nD  >>  sigma_n * (mu_nD/mu_n)
is algebraically  CV_dz / CV_o >> 1, i.e. the ratio of coefficients of
variation is the scale-free diagnostic.
"""
import numpy as np
from scipy import stats
import json

rng = np.random.default_rng(20260605)
NREP = 10000

def ln_params(mu_n, cv):
    s2 = np.log(cv**2 + 1.0)
    mu_e = np.log(mu_n) - s2/2.0
    return mu_e, np.sqrt(s2)

def draw_lognormal(mu_n, cv, n, reps):
    mu_e, s_e = ln_params(mu_n, cv)
    return rng.lognormal(mean=mu_e, sigma=s_e, size=(reps, n))

# ----- check the parametrization reproduces paper's example numbers -----
def report_params(mu_n, sd_n):
    cv = sd_n/mu_n
    mu_e, s_e = ln_params(mu_n, cv)
    return dict(mu_n=mu_n, sd_n=sd_n, cv=round(cv,3),
                mu_e=round(mu_e,3), sigma_e=round(s_e,3))

paper_examples = {
    "control_0.5pm0.5": report_params(0.5, 0.5),
    "disease_1.5pm3.4": report_params(1.5, 3.4),
}

results = {"paper_examples": paper_examples, "NREP": NREP}

# =====================================================================
# 1) SHAPIRO-WILK pass rate (fraction with p>0.05) vs CV and n
#    Higher CV  <-> larger sigma_e <-> stronger non-normality.
# =====================================================================
cv_grid = [0.2, 0.5, 1.0, 1.5, 2.0, 2.27, 3.0]
n_grid  = [10, 20, 30, 50, 100, 200]
sw = {}
for cv in cv_grid:
    sw[cv] = {}
    for n in n_grid:
        reps = 5000
        X = draw_lognormal(1.0, cv, n, reps)
        # Shapiro-Wilk per row
        p = np.array([stats.shapiro(row).pvalue for row in X])
        sw[cv][n] = float(np.mean(p > 0.05))  # "pass" = fail to reject normality
results["shapiro_pass_rate"] = sw

# =====================================================================
# 2) t-TEST CONSEQUENCES
#    (a) Type I error under the null (identical lognormals), raw vs log,
#        with EQUAL and UNEQUAL group sizes (unequal n + heavy tail is
#        where Student's t breaks).
#    (b) Power under a true log-domain shift, raw Welch vs log t-test.
# =====================================================================
def ttest_rates(cv, n1, n2, log_shift, reps=10000, equal_var_student=True):
    """Return (reject_raw_student, reject_raw_welch, reject_log)."""
    mu_e1, s_e = ln_params(1.0, cv)
    # group2: shift median by log_shift in log domain (keeps shape/CV)
    g1 = rng.lognormal(mu_e1, s_e, size=(reps, n1))
    g2 = rng.lognormal(mu_e1 + log_shift, s_e, size=(reps, n2))
    t_s = stats.ttest_ind(g1, g2, axis=1, equal_var=True).pvalue
    t_w = stats.ttest_ind(g1, g2, axis=1, equal_var=False).pvalue
    t_l = stats.ttest_ind(np.log(g1), np.log(g2), axis=1, equal_var=True).pvalue
    return float(np.mean(t_s < .05)), float(np.mean(t_w < .05)), float(np.mean(t_l < .05))

# Type I error (no true difference), unequal n = 30 vs 10
typeI = {}
for cv in [0.2, 0.5, 1.0, 1.5, 2.27, 3.0]:
    s, w, l = ttest_rates(cv, 30, 10, log_shift=0.0)
    typeI[cv] = dict(raw_student=s, raw_welch=w, log=l)
results["typeI_unequal_n_30_10"] = typeI

# Type I error with EQUAL n = 20 vs 20 (robustness baseline)
typeI_eq = {}
for cv in [0.2, 0.5, 1.0, 1.5, 2.27, 3.0]:
    s, w, l = ttest_rates(cv, 20, 20, log_shift=0.0)
    typeI_eq[cv] = dict(raw_student=s, raw_welch=w, log=l)
results["typeI_equal_n_20_20"] = typeI_eq

# Power at a fixed log-domain shift (ln(1.5)=0.405 median ratio), equal n
shift = np.log(1.5)
power = {}
for cv in [0.5, 1.0, 1.5, 2.27, 3.0]:
    pw = {}
    for n in [20, 30, 50, 100]:
        s, w, l = ttest_rates(cv, n, n, log_shift=shift)
        pw[n] = dict(raw_welch=w, log=l)
    power[cv] = pw
results["power_logshift_ln1.5"] = power

# =====================================================================
# 3) CV-RATIO CUTOFF
#    Two groups: control CV_o, disease CV_dz, with a real mean increase.
#    Question: at what observed CV ratio does SW reliably flag, and where
#    does raw inference diverge from log inference?
#    We map: detection probability that AT LEAST ONE group fails SW, and
#    the divergence |reject_raw - reject_log|, as a function of CV ratio.
# =====================================================================
cutoff = {}
cv_o = 1.0
for n in [20, 30, 50]:
    cutoff[n] = {}
    for ratio in [1.0, 1.25, 1.5, 1.75, 2.0, 2.5, 3.0]:
        cv_dz = cv_o * ratio
        # control vs disease, disease mean 1.5x (paper restriction eq 14)
        g_o  = draw_lognormal(1.0, cv_o, n, 5000)
        g_dz = draw_lognormal(1.5, cv_dz, n, 5000)
        sw_o  = np.array([stats.shapiro(r).pvalue for r in g_o])
        sw_dz = np.array([stats.shapiro(r).pvalue for r in g_dz])
        flagged = np.mean((sw_o < .05) | (sw_dz < .05))
        # inference divergence: raw welch vs log on this real comparison
        rw = stats.ttest_ind(g_o, g_dz, axis=1, equal_var=False).pvalue
        lg = stats.ttest_ind(np.log(g_o), np.log(g_dz), axis=1, equal_var=False).pvalue
        cutoff[n][ratio] = dict(sw_flag_rate=float(flagged),
                                power_raw=float(np.mean(rw < .05)),
                                power_log=float(np.mean(lg < .05)))
results["cutoff_by_cvratio"] = cutoff

# =====================================================================
# 4) The paper's exact example at clinical sample sizes
# =====================================================================
ex = {}
for n in [10, 20, 30, 50]:
    g_o  = draw_lognormal(0.5, 1.0,  n, 5000)      # 0.5 +/- 0.5, CV=1.0
    g_dz = draw_lognormal(1.5, 3.4/1.5, n, 5000)   # 1.5 +/- 3.4, CV=2.267
    sw_o  = np.mean(np.array([stats.shapiro(r).pvalue for r in g_o]) > .05)
    sw_dz = np.mean(np.array([stats.shapiro(r).pvalue for r in g_dz]) > .05)
    rw = stats.ttest_ind(g_o, g_dz, axis=1, equal_var=False).pvalue
    lg = stats.ttest_ind(np.log(g_o), np.log(g_dz), axis=1, equal_var=False).pvalue
    ex[n] = dict(sw_pass_control=float(sw_o), sw_pass_disease=float(sw_dz),
                 power_raw=float(np.mean(rw < .05)), power_log=float(np.mean(lg < .05)))
results["paper_example_sims"] = ex

with open("results.json", "w") as f:
    json.dump(results, f, indent=2)
print(json.dumps(paper_examples, indent=2))
print("done")
