# (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).
"""
Which CV-ratio cutoffs (above and below 1) actually matter, and under what conditions?

Heteroscedasticity breaks the pooled (equal-variance) log-t only through the INTERACTION of
the variance imbalance with the sample-size imbalance. We therefore sweep the CV ratio finely
on both sides of 1, for a range of n-allocations, and read off the band of CV ratios within
which the pooled test keeps nominal size. Convention (Behrens-Fisher): with the smaller-n arm
more variable the test is LIBERAL (size > 0.05); with the larger-n arm more variable it is
CONSERVATIVE (size < 0.05); balanced designs are robust at any CV ratio.

Definitions used here:
  - the CV ratio is computed as CV(small-n arm) / CV(large-n arm), so ratio>1 == the small arm
    is the more variable (the dangerous, liberal direction);
  - the larger-n arm is fixed at CV = 1.0 (sigma_e = 0.833); the small arm's CV is varied;
  - a band [c_lo, c_hi] is the CV-ratio interval keeping pooled size inside a tolerance.
Robustness: the mapping CV-ratio -> log-variance ratio is nonlinear, so the band also depends
on the absolute CV level; we re-extract it at large-arm CV = 0.5, 1.0, 2.0.
"""
import numpy as np, json
from scipy import stats
import warnings; warnings.filterwarnings("ignore")

rng = np.random.default_rng(11)
REPS = 12000
ALPHA = 0.05
def se_of_cv(cv): return np.sqrt(np.log(cv**2 + 1.0))

def pooled_welch_size(cv_small, cv_large, n_small, n_large, reps=REPS):
    se_s, se_l = se_of_cv(cv_small), se_of_cv(cv_large)
    gs = rng.lognormal(0.0, se_s, (reps, n_small))   # equal medians -> size
    gl = rng.lognormal(0.0, se_l, (reps, n_large))
    ls, ll = np.log(gs), np.log(gl)
    ms, ml = ls.mean(1), ll.mean(1)
    vs, vl = ls.var(1, ddof=1), ll.var(1, ddof=1)
    diff = ms - ml
    sp2 = ((n_small-1)*vs + (n_large-1)*vl)/(n_small+n_large-2)
    se_p = np.sqrt(sp2*(1.0/n_small + 1.0/n_large))
    p_pool = 2*stats.t.sf(np.abs(diff/se_p), n_small+n_large-2)
    a, b = vs/n_small, vl/n_large
    se_w = np.sqrt(a+b); df_w = (a+b)**2/(a**2/(n_small-1)+b**2/(n_large-1))
    p_welch = 2*stats.t.sf(np.abs(diff/se_w), df_w)
    return float(np.mean(p_pool<ALPHA)), float(np.mean(p_welch<ALPHA))

def cutoffs(cv_large, n_small, m, grid, tol):
    """return (c_lo, c_hi) CV-ratio band keeping pooled size within tol=(lo,hi)."""
    n_large = int(round(n_small*m))
    sizes = np.array([pooled_welch_size(r*cv_large, cv_large, n_small, n_large)[0] for r in grid])
    # size is monotone increasing in r; interpolate crossings
    def cross(target):
        if sizes.min() > target: return grid[0]   # already above -> band edge at/below grid start
        if sizes.max() < target: return grid[-1]  # never reaches -> band extends beyond grid
        return float(np.interp(target, sizes, grid))
    c_lo = cross(tol[0]); c_hi = cross(tol[1])
    return c_lo, c_hi, sizes, n_large

GRID = np.round(np.arange(0.40, 3.01, 0.10), 2)
TOL  = (0.04, 0.06)          # "negligible distortion" band
TOLs = (0.045, 0.055)        # strict
MS = [1.0, 1.5, 2.0, 3.0, 4.0]

print("="*74)
print("MAIN: CV-ratio band [c_lo,c_hi] keeping pooled log-t size in 0.04-0.06")
print("      (CV ratio = small-n arm CV / large-n arm CV;  large-n arm CV = 1.0)")
print("="*74)
print(f"{'n_small/n_large':>16}{'imbalance':>11}{'band 0.04-0.06':>22}{'strict 0.045-0.055':>22}")
out={"main":[]}
for n_small in [10, 20]:
    for m in MS:
        lo,hi,sizes,nl = cutoffs(1.0, n_small, m, GRID, TOL)
        los,his,_,_ = cutoffs(1.0, n_small, m, GRID, TOLs)
        b = "ALL (robust)" if (lo<=GRID[0] and hi>=GRID[-1]) else f"[{lo:.2f}, {hi:.2f}]"
        bs = "ALL" if (los<=GRID[0] and his>=GRID[-1]) else f"[{los:.2f}, {his:.2f}]"
        print(f"{str(n_small)+'/'+str(nl):>16}{m:>10.1f}x{b:>22}{bs:>22}")
        out["main"].append(dict(n_small=n_small,n_large=nl,m=m,c_lo=round(lo,2),c_hi=round(hi,2),
                                c_lo_strict=round(los,2),c_hi_strict=round(his,2),
                                sizes=[round(float(x),3) for x in sizes]))

print("\n"+"="*74)
print("DIRECTION & MAGNITUDE: pooled size at selected CV ratios (n_small=10)")
print("="*74)
print(f"{'CVratio(small/large)':>22}", end="")
for m in MS: print(f"{'m='+str(m):>9}", end="")
print()
for r in [0.5, 0.7, 1.0, 1.5, 2.0, 3.0]:
    print(f"{r:>22.1f}", end="")
    for m in MS:
        nl=int(round(10*m)); s=pooled_welch_size(r*1.0,1.0,10,nl)[0]
        print(f"{s:>9.3f}", end="")
    print()
print("  (Welch held 0.045-0.055 in every one of these cells.)")

print("\n"+"="*74)
print("ROBUSTNESS: band depends on the ABSOLUTE CV level (n_small=10, m=3x)")
print("="*74)
for cvL in [0.5, 1.0, 2.0]:
    lo,hi,_,nl = cutoffs(cvL, 10, 3.0, GRID, TOL)
    print(f"   large-arm CV={cvL}:  safe CV-ratio band [{lo:.2f}, {hi:.2f}]")

# also store the direction table and a fine curve for plotting
fine = np.round(np.arange(0.35,3.21,0.05),2)
curves={}
for m in MS:
    nl=int(round(10*m))
    curves[str(m)] = {"n_large":nl,
        "ratio":[float(x) for x in fine],
        "pool":[round(pooled_welch_size(r*1.0,1.0,10,nl)[0],3) for r in fine],
        "welch":[round(pooled_welch_size(r*1.0,1.0,10,nl)[1],3) for r in fine]}
out["curves"]=curves
json.dump(out, open("cvratio_cutoffs.json","w"))
print("\nsaved cvratio_cutoffs.json")
