"""Numbers for /guides/cohens-d-vs-hedges-g/ (Python 3.13, SciPy 1.18.1).
Summary statistics are the constructed ALT data from /guides/reviewer-asks-to-justify-sample-size/
(6 vs 6 mice), rounded to two decimals as a reader would copy them from a table."""
import json
import math

import numpy as np
from scipy import optimize, special, stats

out = {}


def J_exact(df):
    # Hedges (1981): J = Gamma(df/2) / (sqrt(df/2) * Gamma((df-1)/2)); log-gamma for stability
    return math.exp(special.gammaln(df / 2) - 0.5 * math.log(df / 2) - special.gammaln((df - 1) / 2))


def J_approx(n1, n2):
    return 1 - 3 / (4 * (n1 + n2) - 9)  # Lakens 2013, eq. 4


def ncp_ci(t, df, level=0.95):
    """CI for the noncentrality parameter by inverting the noncentral t CDF (pivot method).
    Solved on |t| and reflected; NaN lower tails (SciPy issue #25470) treated as 0."""
    a = (1 - level) / 2
    s, at = (1 if t >= 0 else -1), abs(t)

    def F(lam):
        v = stats.nct.cdf(at, df, lam)
        return 0.0 if np.isnan(v) else v

    lo = optimize.brentq(lambda L: F(L) - (1 - a), at - 10, at + 10, xtol=1e-14)
    hi = optimize.brentq(lambda L: F(L) - a, at - 10, at + 10, xtol=1e-14)
    return (lo, hi) if s > 0 else (-hi, -lo)


# 1. J by group size (equal groups)
rows = []
for n in [3, 4, 5, 6, 8, 10, 15, 20, 30, 50]:
    df = 2 * n - 2
    j = J_exact(df)
    rows.append({"n_per_group": n, "df": df, "J_exact": j, "J_approx": J_approx(n, n),
                 "g_shrink_pct": 100 * (1 - j)})
out["J_table"] = rows

# 2. Worked example from summary statistics
m1, s1, n1 = 47.17, 7.94, 6  # vehicle
m2, s2, n2 = 58.17, 8.13, 6  # treated
df = n1 + n2 - 2
sp = math.sqrt(((n1 - 1) * s1**2 + (n2 - 1) * s2**2) / df)
d = (m2 - m1) / sp
t_pooled = d / math.sqrt(1 / n1 + 1 / n2)
j = J_exact(df)
g = j * d
lam = ncp_ci(t_pooled, df)
k = math.sqrt(1 / n1 + 1 / n2)
delta_ci = (lam[0] * k, lam[1] * k)
out["example"] = {"diff": m2 - m1, "sd_pooled": sp, "d": d, "J": j, "J_approx": J_approx(n1, n2), "g": g,
                  "t_pooled": t_pooled, "delta_ci": delta_ci,
                  "g_ci_J_scaled_effectsize_style": (j * delta_ci[0], j * delta_ci[1]),
                  "glass_delta_control_sd": (m2 - m1) / s1}

# 3. Paired designs: d_z vs d_av for the same mean change, SD1 = SD2 = 1
paired = []
for r in [0.0, 0.3, 0.5, 0.7, 0.8, 0.9]:
    mdiff, sd1, sd2 = 0.5, 1.0, 1.0
    sdiff = math.sqrt(sd1**2 + sd2**2 - 2 * r * sd1 * sd2)  # Lakens 2013, eq. 8
    paired.append({"r": r, "d_av": mdiff / ((sd1 + sd2) / 2), "d_z": mdiff / sdiff,
                   "ratio_dz_over_dav": (mdiff / sdiff) / (mdiff / ((sd1 + sd2) / 2))})
out["paired"] = paired


# 4. Exact coverage of a J-scaled interval vs the unscaled delta interval (equal groups)
def coverage_scaled(n, delta, level=0.95):
    df = 2 * n - 2
    k = math.sqrt(2 / n)
    a = (1 - level) / 2
    lam_true = delta / k
    lam_prime = (delta / J_exact(df)) / k  # [J*L, J*U] contains delta  <=>  [L, U] contains delta/J
    q_lo = stats.nct.ppf(a, df, lam_prime)
    q_hi = stats.nct.ppf(1 - a, df, lam_prime)
    return float(stats.nct.cdf(q_hi, df, lam_true) - stats.nct.cdf(q_lo, df, lam_true))


out["coverage"] = [{"n_per_group": n, "delta": dl, "coverage_J_scaled": coverage_scaled(n, dl), "coverage_unscaled": 0.95}
                   for n, dl in [(3, 1.0), (3, 2.0), (6, 1.0), (6, 2.0), (6, 3.0), (10, 1.0), (20, 1.0)]]


def clean(o):
    if isinstance(o, dict):
        return {k: clean(v) for k, v in o.items()}
    if isinstance(o, (list, tuple)):
        return [clean(v) for v in o]
    if isinstance(o, float):
        return float(f"{o:.12g}")
    return o


with open("example_output.json", "w") as fh:
    json.dump(clean(out), fh, indent=1)
print(json.dumps(clean(out), indent=1))