# (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).
"""Compare normality tests vs sample size on log-normal data.
Tests: Shapiro-Wilk (SW), Anderson-Darling (AD), Lilliefors/KS (LF),
D'Agostino-Pearson omnibus (DP), Jarque-Bera (JB).
Outputs: power (reject rate) to detect log-normal departures by CV and n,
and size (type I) under true normality by n.
"""
import numpy as np, json, warnings
from scipy import stats
from statsmodels.stats.diagnostic import lilliefors
warnings.filterwarnings("ignore")
rng = np.random.default_rng(424242)
REPS = 2000
ALPHA = 0.05

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

def reject_rates(X):
    """X: reps x n array. Return dict of rejection rate per test."""
    reps = X.shape[0]
    sw = np.mean([stats.shapiro(r).pvalue < ALPHA for r in X])
    # Anderson-Darling: reject if stat >= 5% critical value (index 2)
    ad_flags = []
    for r in X:
        res = stats.anderson(r, "norm"); ad_flags.append(res.statistic >= res.critical_values[2])
    ad = float(np.mean(ad_flags))
    lf = np.mean([lilliefors(r, dist="norm")[1] < ALPHA for r in X])
    dp = float(np.mean(stats.normaltest(X, axis=1).pvalue < ALPHA))     # D'Agostino-Pearson
    # Jarque-Bera vectorized: JB = n/6*(S^2 + (K^2)/4), K=excess kurtosis; p from chi2(2)
    n = X.shape[1]
    S = stats.skew(X, axis=1); K = stats.kurtosis(X, axis=1)  # excess kurtosis
    jb_stat = n / 6.0 * (S**2 + (K**2) / 4.0)
    jb = float(np.mean(stats.chi2.sf(jb_stat, 2) < ALPHA))
    return {"SW": float(sw), "AD": ad, "LF": float(lf), "DP": dp, "JB": jb}

n_grid = [10, 20, 30, 50, 100, 200]
cv_grid = [0.5, 1.0, 2.0]

power = {}   # power[cv][n] = {test: rate}
for cv in cv_grid:
    power[cv] = {}
    for n in n_grid:
        mu_e, s_e = ln(1.0, cv)
        X = rng.lognormal(mu_e, s_e, size=(REPS, n))
        power[cv][n] = reject_rates(X)

size = {}    # size[n] = {test: rate} under true normal
for n in n_grid:
    X = rng.normal(0, 1, size=(REPS, n))
    size[n] = reject_rates(X)

out = {"power": power, "size": size, "REPS": REPS, "n_grid": n_grid, "cv_grid": cv_grid}
json.dump(out, open("normtests.json", "w"), indent=2)

# console summary
tests = ["SW", "AD", "LF", "DP", "JB"]
print("POWER to detect log-normal (reject rate at alpha=0.05):")
for cv in cv_grid:
    print(f"\nCV={cv}")
    print("  n   " + "  ".join(f"{t:>5}" for t in tests))
    for n in n_grid:
        print(f"  {n:>3} " + "  ".join(f"{power[cv][n][t]:>5.2f}" for t in tests))
print("\nSIZE under true normality (should be ~0.05):")
print("  n   " + "  ".join(f"{t:>5}" for t in tests))
for n in n_grid:
    print(f"  {n:>3} " + "  ".join(f"{size[n][t]:>5.2f}" for t in tests))
print("\ndone")
