# DASS Analysis Clinic Case 005; 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))

effect=.35; m=25
n=2*(stats.norm.ppf(.975)+stats.norm.ppf(.8))**2/effect**2
print('Normal approximation only; individual n per arm:',n)
for icc in [.01,.03,.05,.10]:
 de=1+(m-1)*icc; k=int(np.ceil(n*de/m)); print(f'ICC={icc:.2f}; DE={de:.2f}; clusters/arm={k}; people/arm={k*m}')
print('20 retained per cluster, ICC .05:',int(np.ceil(n*(1+19*.05)/20)))
print('12 analyzable clusters and 10% cluster loss:',int(np.ceil(12/.9)))
assert [int(np.ceil(n*(1+24*r)/25)) for r in [.01,.03,.05,.10]]==[7,9,12,18]