"""Numbers for /guides/how-many-biological-replicates/ (Python 3.13, SciPy 1.18.1, NumPy).
Two-sample t-test, equal n, two-sided alpha = 0.05, power from the noncentral t distribution
(same method as /tools/sample-size-calculator/). Run from this directory: python example.py"""
import json
import platform
import numpy as np
import scipy
from scipy import stats

ALPHA = 0.05


def power(n, d, alpha=ALPHA):
    df = 2 * n - 2
    ncp = d * np.sqrt(n / 2)
    tc = stats.t.ppf(1 - alpha / 2, df)
    return stats.nct.sf(tc, df, ncp) + stats.nct.cdf(-tc, df, ncp)


def n_for(d, target, alpha=ALPHA):
    n = 2
    while power(n, d, alpha) < target:
        n += 1
    return n


out = {"versions": {"python": platform.python_version(), "scipy": scipy.__version__, "numpy": np.__version__}}

# 1. Table: n per group for 80% and 90% power
table = []
for d in [0.5, 0.8, 1.0, 1.5, 2.0]:
    row = {"d": d}
    for target in (0.8, 0.9):
        n = n_for(d, target)
        row[f"n_{int(target * 100)}"] = n
        row[f"achieved_{int(target * 100)}"] = round(float(power(n, d)), 4)
    table.append(row)
out["table"] = table

# 2. Worked planning example: difference 4 g, SD 3 g between animals (d = 1.33)
diff, sd = 4.0, 3.0
d = diff / sd
n80 = n_for(d, 0.8)
out["example"] = {"diff": diff, "sd": sd, "d": round(d, 4), "n80": n80, "achieved": round(float(power(n80, d)), 4),
                  "n80_with_10pct_loss": int(np.ceil(n80 / 0.9)),
                  "sd_4_n80": n_for(diff / 4.0, 0.8),
                  "sd_4_power_at_n80": round(float(power(n80, diff / 4.0)), 4)}

# 3. Resource equation: E = N - k between 10 and 20
res = []
for k in (2, 3, 4):
    lo, hi = 10 + k, 20 + k
    res.append({"groups": k, "total_min": lo, "total_max": hi,
                "per_group_min": int(np.ceil(lo / k)), "per_group_max": hi // k})
out["resource_equation"] = res

# 4. Pilot-SD simulation: true SD 1, true difference 1 (d = 1). A pilot of 5 animals per group gives
# a pooled SD on 8 df; the main study is sized for 80% power with that SD. What power does it really have?
rng = np.random.default_rng(20261110)
R = 20000
pilot_n = 5
df_p = 2 * pilot_n - 2
true_d = 1.0
s = np.sqrt(rng.chisquare(df_p, R) / df_p)  # pooled pilot SD / true SD
n_cache = {}
p_cache = {}


def n_cached(dd):
    key = round(float(dd), 3)
    if key not in n_cache:
        n_cache[key] = n_for(key, 0.8)
    return n_cache[key]


def p_cached(n):
    if n not in p_cache:
        p_cache[n] = float(power(n, true_d))
    return p_cache[n]


n_plain = np.array([n_cached(true_d / si) for si in s])
pow_plain = np.array([p_cached(int(n)) for n in n_plain])
# planning with the upper one-sided 80% confidence limit of the SD (Browne 1995)
ucl = np.sqrt(df_p / stats.chi2.ppf(0.2, df_p))
n_ucl = np.array([n_cached(true_d / (si * ucl)) for si in s])
pow_ucl = np.array([p_cached(int(n)) for n in n_ucl])
out["pilot_simulation"] = {
    "replicates": R, "seed": 20261110, "pilot_per_group": pilot_n, "true_d": true_d,
    "n_needed_true": n_for(true_d, 0.8),
    "share_pilot_sd_below_true": round(float(np.mean(s < 1)), 4),
    "plain": {"share_power_below_80": round(float(np.mean(pow_plain < 0.8)), 4),
              "share_power_below_50": round(float(np.mean(pow_plain < 0.5)), 4),
              "median_power": round(float(np.median(pow_plain)), 4),
              "median_n": int(np.median(n_plain))},
    "ucl80_multiplier": round(float(ucl), 4),
    "ucl80": {"share_power_below_80": round(float(np.mean(pow_ucl < 0.8)), 4),
              "share_power_below_50": round(float(np.mean(pow_ucl < 0.5)), 4),
              "median_power": round(float(np.median(pow_ucl)), 4),
              "median_n": int(np.median(n_ucl))},
}


# 5. Sensitivity: smallest d detectable with 80% power for a feasible n
def d_detect(n, target=0.8):
    lo, hi = 0.01, 10.0
    for _ in range(100):
        mid = (lo + hi) / 2
        if power(n, mid) < target:
            lo = mid
        else:
            hi = mid
    return hi


out["sensitivity"] = {str(n): round(d_detect(n), 3) for n in (5, 6, 8, 10, 12)}

with open("example_output.json", "w") as f:
    json.dump(out, f, indent=2)
print(json.dumps(out, indent=2))