# DASS Analysis Clinic Case 003; Python dependencies: numpy, scipy.
from pathlib import Path
import numpy as np
from scipy import stats, optimize
HERE = Path(__file__).resolve().parent

def ols(X, y):
    b = np.linalg.lstsq(X, y, rcond=None)[0]
    residual = y - X @ b
    df = len(y) - X.shape[1]
    cov = (residual @ residual / df) * np.linalg.inv(X.T @ X)
    se = np.sqrt(np.diag(cov))
    ci = np.column_stack((b-stats.t.ppf(.975, df)*se,b+stats.t.ppf(.975,df)*se))
    return b, cov, ci

def power(n, d):
    # Matches R power.t.test(strict=FALSE): rejection tail in effect direction.
    return stats.nct.sf(stats.t.ppf(.975,2*n-2),2*n-2,d*np.sqrt(n/2))

pilots=np.genfromtxt(HERE/'case-003-shared-data.csv',delimiter=',',names=True)
def planned(d):
 if d<=.05:return np.inf
 return np.ceil(optimize.brentq(lambda n:power(n,d)-.8,2,1e6))
naive=np.array([planned(d) for d in pilots['d_hat']]); safe=np.array([planned(d) for d in pilots['d_safe']])
go=naive<=150; go_safe=safe<=150
powers=power(naive[go],.4)
print(f'True d=.40: {planned(.4):.0f} per group; pilot 95% range={np.quantile(pilots["d_hat"],[.025,.975])}')
print(f'Go-ahead={go.mean():.3f}; mean selected d={pilots["d_hat"][go].mean():.3f}; median power={np.median(powers):.3f}; below .80={np.mean(powers<.8):.3f}')
print(f'SESOI .35: {planned(.35):.0f} per group; safeguard go-ahead={go_safe.mean():.3f}; median power={np.median(power(safe[go_safe],.4)):.3f}')
assert planned(.4)==100 and planned(.35)==130