"""False-positive rates of four testing policies, and sample size for one- vs two-sided tests,
for /guides/one-vs-two-tailed-test/. Python 3.13, SciPy 1.18.1.
Run from this folder: python simulation.py > simulation_output.txt  (also writes simulation_output.json)

Part 1: no true difference; 8 animals per group; normal data; Welch's t-test; 100,000 experiments; seed 20261110.
Part 2: animals per group for a two-sample Student t-test (the sample size calculator's model), found by exact
search with the noncentral t distribution."""
import json

import numpy as np
from scipy import stats

rng = np.random.default_rng(20261110)
N_SIM, n, alpha = 100_000, 8, 0.05
a = rng.normal(0.0, 1.0, size=(N_SIM, n))   # control
b = rng.normal(0.0, 1.0, size=(N_SIM, n))   # treated; same distribution, so every "significant" result is false

va, vb = a.var(axis=1, ddof=1) / n, b.var(axis=1, ddof=1) / n
t = (b.mean(axis=1) - a.mean(axis=1)) / np.sqrt(va + vb)
df = (va + vb) ** 2 / (va ** 2 / (n - 1) + vb ** 2 / (n - 1))
p_two = 2 * stats.t.sf(np.abs(t), df)
p_less = stats.t.cdf(t, df)         # predicted direction: treated lower
p_greater = stats.t.sf(t, df)

policies = {
    "two_sided": p_two < alpha,
    "one_sided_direction_fixed_in_advance": p_less < alpha,
    "switch_to_one_sided_if_it_helps": (p_two < alpha) | (np.minimum(p_less, p_greater) < alpha),
    "one_sided_planned_but_wrong_direction_also_reported": (p_less < alpha) | (p_two < alpha),
    "one_sided_at_0.025": p_less < alpha / 2,
}
rates = {k: float(v.mean()) for k, v in policies.items()}


def power_student(nn, d, alpha, sides):
    dfree = 2 * nn - 2
    ncp = d * np.sqrt(nn / 2)
    if sides == 2:
        crit = stats.t.ppf(1 - alpha / 2, dfree)
        return stats.nct.sf(crit, dfree, ncp) + stats.nct.cdf(-crit, dfree, ncp)
    crit = stats.t.ppf(1 - alpha, dfree)
    return stats.nct.sf(crit, dfree, ncp)


def n_needed(d, alpha, target, sides):
    nn = 2
    while power_student(nn, d, alpha, sides) < target:
        nn += 1
    return nn


sizes = []
for target in (0.80, 0.90):
    for d in (0.5, 0.8, 1.0, 1.5, 2.0):
        two = n_needed(d, 0.05, target, 2)
        one = n_needed(d, 0.05, target, 1)
        sizes.append({"power": target, "d": d, "n_two_sided": two, "n_one_sided": one,
                      "saving_percent": round(100 * (two - one) / two, 1)})

out = {"part1": {"n_per_group": n, "simulations": N_SIM, "seed": 20261110, "rates": rates},
       "part2": sizes}
print("False-positive rate with no true difference (8 per group, 100,000 experiments, seed 20261110):")
for k, v in rates.items():
    print(f"  {k:55s} {v:.4f}")
print("Animals per group, two-sample Student t-test, alpha 0.05:")
for s in sizes:
    print(f"  power {s['power']:.2f} d {s['d']:.1f}: two-sided {s['n_two_sided']:3d}  one-sided {s['n_one_sided']:3d}  saving {s['saving_percent']}%")
with open("simulation_output.json", "w") as f:
    json.dump(out, f, indent=2)