# (c) 2026 William J. Dwyer. MIT License. Cite the methodology paper and archived software
# (Zenodo DOI 10.5281/zenodo.20709963).
"""Verification of the GAMMA rescue pathway for the coefficient-of-variation screen.

When the three-way model check (cvscreen_core.model_check) routes a flagged outcome to "gamma"
rather than "lognormal", the rescue is NOT the log-scale t-test (the logarithm does not normalize a
gamma variate) but a gamma generalized linear model with a log link. This script establishes, by
Monte-Carlo simulation, the three claims the manuscript's gamma-rescue subsection makes:

  (1) CALIBRATION.  On gamma data with mean-proportional dispersion, the raw-scale Welch t-test is
      mis-calibrated under sample-size imbalance and the log-scale t-test is only approximately
      valid, whereas the gamma GLM (log link) holds the nominal level.
  (2) POWER.  For a true multiplicative shift in the MEAN, the gamma GLM is at least as powerful as
      the log-scale t-test and more powerful than the raw-scale t-test.
  (3) ESTIMAND.  The gamma GLM targets the ratio of ARITHMETIC means (exp(beta)); the log-scale
      t-test silently targets the ratio of GEOMETRIC means.  When the two groups differ in SHAPE,
      these diverge: the GLM recovers the mean ratio without bias while the log-t reports the
      (different) geometric-mean ratio -- a change of estimand, not merely of efficiency.

The two-group gamma GLM with a log link is fitted in closed form: for a saturated two-group model
the fitted group means are the sample means, and the Wald variance of log(mean) is the Pearson
dispersion over n.  Allowing a SEPARATE dispersion per group is the gamma analogue of Welch's
separate-variances t-test.  The closed form is cross-checked against statsmodels when available.
"""
import sys; sys.path.insert(0, "/sessions/optimistic-keen-clarke/mnt/m_series")
import numpy as np
from scipy import stats
from scipy.special import polygamma, digamma
import cvscreen_core as cc

# ----------------------------------------------------------------------------- helpers
def gamma_samples(rng, n, mean, k):
    """n draws from Gamma(shape=k, scale=mean/k): E[X]=mean, CV=1/sqrt(k)."""
    return rng.gamma(k, mean / k, size=n)

def welch_t_vec(A, B):
    """Vectorized two-sided Welch t-test p-values over rows (each row is one replicate)."""
    na, nb = A.shape[1], B.shape[1]
    ma, mb = A.mean(1), B.mean(1); va, vb = A.var(1, ddof=1), B.var(1, ddof=1)
    se2a, se2b = va / na, vb / nb
    t = (mb - ma) / np.sqrt(se2a + se2b)
    df = (se2a + se2b) ** 2 / (se2a ** 2 / (na - 1) + se2b ** 2 / (nb - 1))
    return 2 * stats.t.sf(np.abs(t), df)

def gamma_glm_vec(A, B, separate_dispersion=True):
    """Vectorized closed-form two-group gamma GLM (log link) over rows.
    Returns (log_mean_ratio, two_sided_p).  exp(log_mean_ratio) is the ratio of ARITHMETIC means;
    Pearson dispersion phi_g = mean(((x-mu_g)/mu_g)^2) estimates CV^2; Var(log mu_g)=phi_g/n_g.
    separate_dispersion=True is the Welch analogue (separate shape per arm)."""
    na, nb = A.shape[1], B.shape[1]
    ma, mb = A.mean(1), B.mean(1)
    pa = (((A - ma[:, None]) / ma[:, None]) ** 2).mean(1)   # Pearson dispersion (sample CV^2), arm a
    pb = (((B - mb[:, None]) / mb[:, None]) ** 2).mean(1)
    if separate_dispersion:
        va, vb = pa / na, pb / nb
    else:
        phi = (na * pa + nb * pb) / (na + nb)
        va, vb = phi / na, phi / nb
    lr = np.log(mb) - np.log(ma)
    z = lr / np.sqrt(va + vb)
    return lr, 2 * stats.norm.sf(np.abs(z))

def gamma_glm_loglink(a, b, separate_dispersion=True):
    """Scalar wrapper used only by the statsmodels cross-check."""
    lr, p = gamma_glm_vec(a[None, :], b[None, :], separate_dispersion)
    return lr[0], None, p[0]

def _gamma_unit_deviance(X, mu):
    """Summed gamma deviance per row for data rows X (reps,n) and fitted mean mu (reps,1)."""
    return 2.0 * np.sum(-np.log(X / mu) + (X - mu) / mu, axis=1)

def gamma_glm_F_vec(A, B):
    """Vectorized analysis-of-deviance F-test for the two-group gamma GLM (log link) -- the standard
    small-sample inference (R's anova(., test='F'), statsmodels GLM F-test).  Compares the null
    (common mean) to the full model (separate means); F = (D0 - D1)/phi ~ F(1, n-2)."""
    na, nb = A.shape[1], B.shape[1]; n = na + nb
    ma, mb = A.mean(1)[:, None], B.mean(1)[:, None]
    grand = ((A.sum(1) + B.sum(1)) / n)[:, None]
    D1 = _gamma_unit_deviance(A, ma) + _gamma_unit_deviance(B, mb)        # separate means
    D0 = _gamma_unit_deviance(A, grand) + _gamma_unit_deviance(B, grand)  # common mean
    pa = (((A - ma) / ma) ** 2).sum(1); pb = (((B - mb) / mb) ** 2).sum(1)
    phi = (pa + pb) / (n - 2)                                            # Pearson dispersion
    F = (D0 - D1) / phi
    return stats.f.sf(F, 1, n - 2)

def geom_mean_ratio_truth(k1, k2):
    """True ratio of geometric means for two gammas with EQUAL arithmetic mean but shapes k1,k2.
    GM_g = (mu/k_g) exp(psi(k_g)); with mu equal, GM2/GM1 = (k1/k2) exp(psi(k2)-psi(k1))."""
    return (k1 / k2) * np.exp(digamma(k2) - digamma(k1))

# ----------------------------------------------------------------------------- cross-check vs statsmodels
def crosscheck_statsmodels():
    try:
        import statsmodels.api as sm
    except Exception:
        return "statsmodels unavailable - closed form not cross-checked"
    rng = np.random.default_rng(7)
    a = gamma_samples(rng, 40, 10.0, 2.0); b = gamma_samples(rng, 60, 14.0, 2.0)
    y = np.concatenate([a, b]); X = sm.add_constant(np.concatenate([np.zeros(40), np.ones(60)]))
    m = sm.GLM(y, X, family=sm.families.Gamma(sm.families.links.Log())).fit()
    lr_sm, p_sm = m.params[1], m.pvalues[1]
    lr_cf, z_cf, p_cf = gamma_glm_loglink(a, b, separate_dispersion=False)  # statsmodels uses common phi
    return (f"closed-form vs statsmodels: log-ratio {lr_cf:.4f} vs {lr_sm:.4f} "
            f"(|d|={abs(lr_cf-lr_sm):.4f}); both estimate the arithmetic-mean ratio")

# ----------------------------------------------------------------------------- simulation
def run(reps=6000, seed=20260624):
    rng = np.random.default_rng(seed)
    print("== three-way check routes gamma data to the gamma GLM ==")
    g = rng.gamma(1.0, 5.0, size=6000)                       # Gamma(k=1) -> CV=1
    cv = g.std(ddof=1) / g.mean(); s_log2 = np.log(g).var(ddof=1)
    md = cc.model_discrimination(cv, s_log2)
    print(f"   gamma(k=1) sample: CV={cv:.3f}  s_y^2={s_log2:.3f}  "
          f"lognormal ln(1+CV^2)={md['lognormal_log_variance']:.3f}  "
          f"gamma psi'(1/CV^2)={md['gamma_log_variance']:.3f}  ->  verdict='{md['model']}'  "
          f"rescue='{md['recommended']}'")
    print("  ", crosscheck_statsmodels())

    CVs = [0.5, 1.0, 2.0]
    allocs = [(20, 20), (30, 10)]
    print("\n== (1) calibration  &  (2) power on gamma data (rejection rate at alpha=0.05) ==")
    print(f"   {'CV':>4} {'n1,n2':>7} {'R':>4} | {'rawWelch':>9} {'logWelch':>9} {'gammaGLM':>9}")
    for cvv in CVs:
        k = 1.0 / cvv ** 2
        for (n1, n2) in allocs:
            for R in (1.0, 2.0):                              # R = true ratio of ARITHMETIC means
                A = rng.gamma(k, 10.0 / k, size=(reps, n1))
                B = rng.gamma(k, (10.0 * R) / k, size=(reps, n2))   # equal shape (constant CV)
                p_raw = welch_t_vec(A, B)
                p_log = welch_t_vec(np.log(A), np.log(B))
                p_glm = gamma_glm_F_vec(A, B)                  # analysis-of-deviance F-test
                rej = [(p < 0.05).mean() for p in (p_raw, p_log, p_glm)]
                tag = "size" if R == 1.0 else "powr"
                print(f"   {cvv:>4} {f'{n1},{n2}':>7} {R:>4.0f} | "
                      f"{rej[0]:>9.3f} {rej[1]:>9.3f} {rej[2]:>9.3f}   ({tag})")

    print("\n== (3) estimand divergence: groups differ in SHAPE, equal arithmetic mean (R=1) ==")
    print("   gamma GLM should report mean-ratio ~1 and stay calibrated; the log-t targets the")
    print("   geometric-mean ratio, which is genuinely != 1 when shapes differ.")
    k1, k2 = 4.0, 1.0                                          # CV 0.5 vs CV 1.0
    gm_truth = geom_mean_ratio_truth(k1, k2)
    n1 = n2 = 40
    A = rng.gamma(k1, 10.0 / k1, size=(reps, n1))
    B = rng.gamma(k2, 10.0 / k2, size=(reps, n2))             # SAME arithmetic mean, different shape
    glm_lr, _ = gamma_glm_vec(A, B, separate_dispersion=True)
    glm_p = gamma_glm_F_vec(A, B)
    logt_lr = np.log(B).mean(1) - np.log(A).mean(1)
    logt_p = welch_t_vec(np.log(A), np.log(B))
    glm_ratio = np.exp(glm_lr.mean()); logt_ratio = np.exp(logt_lr.mean())
    print(f"   true arithmetic-mean ratio   = 1.000  (the estimand of interest)")
    print(f"   true geometric-mean ratio    = {gm_truth:.3f}  (what the log-t targets)")
    print(f"   gamma GLM   mean exp(beta)   = {glm_ratio:.3f}   rejection@0.05 = {(glm_p<0.05).mean():.3f}")
    print(f"   log-t       mean exp(dbar)   = {logt_ratio:.3f}   rejection@0.05 = {(logt_p<0.05).mean():.3f}")
    print(f"   => the log-t rejects the equal-MEAN null ~{(logt_p<0.05).mean()*100:.0f}% of the time")
    print(f"      because it is answering the equal-GEOMETRIC-mean question; the GLM stays near {0.05:.2f}.")

if __name__ == "__main__":
    run()
