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

# Specialized ML for this complete balanced Gaussian crossed-intercept design.
# Not a general mixed-model package, and not lme4's convergence/singularity test.
d=np.genfromtxt(HERE/'case-006-shared-data.csv',delimiter=',',names=True)
y=d['y']; x=d['x']; subject=d['subject']; item=d['item']
subjects=np.unique(subject); items=np.unique(item); ns=len(subjects); ni=len(items)
assert len(y)==ns*ni and len(set(zip(subject,item)))==ns*ni
assert all(len(np.unique(x[subject==s]))==1 for s in subjects)
X=np.column_stack((np.ones(len(y)),x)); beta=np.linalg.lstsq(X,y,rcond=None)[0]
r=y-X@beta
grand=np.full(len(y),r.mean())
s=np.array([r[subject==g].mean() for g in subject])-grand
i=np.array([r[item==g].mean() for g in item])-grand
e=r-grand-s-i
ss=np.array([e@e,s@s,i@i,grand@grand]); dims=np.array([(ns-1)*(ni-1),ns-1,ni-1,1])
def objective(logvar):
 vs,vi,ve=np.exp(logvar)
 eigen=np.array([ve,ve+ni*vs,ve+ns*vi,ve+ni*vs+ns*vi])
 return .5*np.sum(dims*np.log(eigen)+ss/eigen)
a=optimize.minimize(objective,np.zeros(3),method='L-BFGS-B',bounds=[(-16,8)]*3)
b=optimize.minimize(objective,np.zeros(3),method='Powell',bounds=[(-16,8)]*3)
print('Within-subject x values: 1; unsupported subject slope is confounded with intercept.')
print('Supported treatment coefficient:',beta[1]); print('Variance components subject/item/residual:',np.exp(a.x))
print('Optimizer success:',a.success,b.success,'objective difference:',abs(a.fun-b.fun))
assert a.success and b.success and abs(a.fun-b.fun)<1e-4 and np.isclose(beta[1],.7950,atol=.00005)