#!/usr/bin/env python3
"""LAB_EXPERIMENT_02 power calculations (stdlib + math only; scipy not required).

ALL effect sizes below are ASSUMPTIONS for illustration. They are not estimates
from data. Run: python3 LAB02_power_calc.py > LAB02_power_calc_output.txt
"""
import math
import random
from statistics import NormalDist

N = NormalDist()
z = N.inv_cdf


def _betacf(a, b, x):
    tiny = 1e-300
    qab, qap, qam = a + b, a + 1, a - 1
    c, d = 1.0, 1 - qab * x / qap
    d = tiny if abs(d) < tiny else d
    d = 1 / d
    h = d
    for m in range(1, 300):
        m2 = 2 * m
        aa = m * (b - m) * x / ((qam + m2) * (a + m2))
        d = 1 + aa * d; d = tiny if abs(d) < tiny else d
        c = 1 + aa / c; c = tiny if abs(c) < tiny else c
        d = 1 / d; h *= d * c
        aa = -(a + m) * (qab + m) * x / ((a + m2) * (qap + m2))
        d = 1 + aa * d; d = tiny if abs(d) < tiny else d
        c = 1 + aa / c; c = tiny if abs(c) < tiny else c
        d = 1 / d
        de = d * c
        h *= de
        if abs(de - 1) < 1e-14:
            break
    return h


def betai(a, b, x):
    if x <= 0: return 0.0
    if x >= 1: return 1.0
    bt = math.exp(math.lgamma(a + b) - math.lgamma(a) - math.lgamma(b)
                  + a * math.log(x) + b * math.log(1 - x))
    if x < (a + 1) / (a + b + 2):
        return bt * _betacf(a, b, x) / a
    return 1 - bt * _betacf(b, a, 1 - x) / b


def t_cdf(t, df):
    x = df / (df + t * t)
    p = 0.5 * betai(df / 2, 0.5, x)
    return 1 - p if t > 0 else p


def t_ppf(p, df):
    lo, hi = -50.0, 50.0
    for _ in range(200):
        mid = (lo + hi) / 2
        if t_cdf(mid, df) < p: lo = mid
        else: hi = mid
    return (lo + hi) / 2


def n_per_arm_two_sample_normal(d, alpha, power):
    """n = 2 (z_{1-a/2} + z_{power})^2 / d^2 (large-sample), +z_{1-a/2}^2/4 (Guenther t-correction)."""
    za, zb = z(1 - alpha / 2), z(power)
    n0 = 2 * (za + zb) ** 2 / d ** 2
    return n0, n0 + za ** 2 / 4


def mc_power_two_sample_t(n, d, alpha, sims=20000, seed=20261111):
    rng = random.Random(seed)
    df = 2 * n - 2
    tcrit = t_ppf(1 - alpha / 2, df)
    hits = 0
    for _ in range(sims):
        a = [rng.gauss(0, 1) for _ in range(n)]
        b = [rng.gauss(d, 1) for _ in range(n)]
        ma, mb = sum(a) / n, sum(b) / n
        va = sum((x - ma) ** 2 for x in a) / (n - 1)
        vb = sum((x - mb) ** 2 for x in b) / (n - 1)
        se = math.sqrt((va + vb) / n)
        if abs((mb - ma) / se) > tcrit:
            hits += 1
    return hits / sims


def n_per_arm_two_prop(p1, p2, alpha, power):
    """n = (z_{1-a/2}*sqrt(2 pbar qbar) + z_pow*sqrt(p1q1+p2q2))^2 / (p1-p2)^2"""
    pbar = (p1 + p2) / 2
    za, zb = z(1 - alpha / 2), z(power)
    num = (za * math.sqrt(2 * pbar * (1 - pbar)) + zb * math.sqrt(p1 * (1 - p1) + p2 * (1 - p2))) ** 2
    return num / (p1 - p2) ** 2


def n_paired(dz, alpha, power):
    """n pairs = (z_{1-a/2}+z_pow)^2 / dz^2 + z_{1-a/2}^2/2 (approx)"""
    za, zb = z(1 - alpha / 2), z(power)
    return (za + zb) ** 2 / dz ** 2 + za ** 2 / 2


def main():
    print("LAB_EXPERIMENT_02 power calculations -- every effect size is an ASSUMPTION, not a result")
    print("Python", __import__("sys").version.split()[0], "| scipy not used | seed 20261111\n")

    print("A. Two-sample t-test, two-sided, ASSUMPTION d=0.5, alpha=0.05, power=0.80")
    n0, n1 = n_per_arm_two_sample_normal(0.5, 0.05, 0.80)
    print(f"  z_(0.975)={z(0.975):.4f}  z_(0.80)={z(0.80):.4f}")
    print(f"  large-sample n/arm = 2(za+zb)^2/d^2 = {n0:.2f} -> {math.ceil(n0)}")
    print(f"  with t-correction  n/arm = {n1:.2f} -> {math.ceil(n1)}")
    nn = math.ceil(n1)
    for n in (nn - 1, nn, nn + 1):
        print(f"  Monte-Carlo power at n={n}/arm (20000 sims): {mc_power_two_sample_t(n, 0.5, 0.05):.4f}")

    print("\nB. Same, d sweep (alpha=0.05, power=0.80), n/arm (t-corrected, rounded up)")
    for d in (0.2, 0.3, 0.4, 0.5, 0.6, 0.8):
        print(f"  d={d}: {math.ceil(n_per_arm_two_sample_normal(d, 0.05, 0.80)[1])}")

    print("\nC. Multiplicity: Bonferroni-equivalent alpha for a primary family of m tests")
    for m in (1, 2, 3, 4):
        a = 0.05 / m
        print(f"  m={m}: alpha_per_test={a:.4f}; d=0.5 power 0.80 -> n/arm "
              f"{math.ceil(n_per_arm_two_sample_normal(0.5, a, 0.80)[1])}  "
              "(Holm is uniformly at least as powerful; this is the conservative planning bound)")

    print("\nD. Two proportions (task-level rates), ASSUMPTION: unsupported-agreement rate 0.30 (control) vs 0.20 (intervention)")
    for pw in (0.80, 0.90):
        print(f"  alpha=0.05 power={pw}: n/arm = {math.ceil(n_per_arm_two_prop(0.30, 0.20, 0.05, pw))}")
    print("  alpha=0.0125 (Bonferroni m=4) power=0.80: n/arm = "
          f"{math.ceil(n_per_arm_two_prop(0.30, 0.20, 0.0125, 0.80))}")

    print("\nE. Clustered runs (repeated runs within task): design effect DEFF = 1+(m-1)*ICC")
    base = math.ceil(n_per_arm_two_prop(0.30, 0.20, 0.05, 0.80))
    for icc in (0.05, 0.10, 0.20):
        for m in (5, 10):
            deff = 1 + (m - 1) * icc
            print(f"  ASSUMPTION ICC={icc}, runs/task m={m}: DEFF={deff:.2f}; runs/arm={math.ceil(base*deff)} "
                  f"(~{math.ceil(base*deff/m)} tasks/arm)")

    print("\nF. Paired / within-participant (e.g. pre/post self-report), ASSUMPTION dz=0.4, alpha=0.05, power=0.80")
    print(f"  pairs needed = {math.ceil(n_paired(0.4, 0.05, 0.80))}")

    print("\nG. Attrition inflation: n_enrol = n / (1 - a); ASSUMPTION a=0.20 applied to A")
    print(f"  {nn} / 0.8 = {math.ceil(nn/0.8)} per arm")


if __name__ == "__main__":
    main()