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

x=np.repeat([-1.,1.],4); w=np.tile([-1.,-1.,1.,1.],2); e=np.tile([-1.,1.],4)
z=x+w; y=-x+2*z+e
naive=ols(np.column_stack((np.ones(8),x)),y)
adjusted=ols(np.column_stack((np.ones(8),x,z)),y)
vif=1/(1-np.corrcoef(x,z)[0,1]**2)
print('Unadjusted x / CI:',naive[0][1],naive[2][1]); print('Adjusted x / CI:',adjusted[0][1],adjusted[2][1]);print('Adjusted z:',adjusted[0][2],'VIF:',vif)
assert np.isclose(naive[0][1],1) and np.isclose(adjusted[0][1],-1) and np.isclose(vif,2)