import numpy as np
import matplotlib.pyplot as plt 
import scipy.stats as st 

np.random.seed(42)

############################################
#       Single Control Variate Case        #
############################################

N = 10**6
sqrt2 = np.sqrt(2) # precomputing factors to make computations faster

print('computing RVs...')
U = st.uniform.rvs(size=(N,2))
Z = [4*int( (u[0]**2+u[1]**2)<=1 ) for u in U]
Y1 = [int( (u[0]+u[1])<=1 ) for u in U]
Y2 = [int( (u[0]+u[1])>=sqrt2 ) for u in U]
Y3 = [(u[0]+u[1]-1)*int( (u[0]+u[1])>1 and (u[0]+u[1])<sqrt2) for u in U]
Y = [Y1,Y2,Y3]

print('computing means..')
meanZ = np.mean(Z)
meanY1 = np.mean(Y1)
meanY2 = np.mean(Y2)
meanY3 = np.mean(Y3)
meanY = [meanY1,meanY2,meanY3]

print('computing variances')
varZ = np.var(Z,ddof=1)
varY1 = np.var(Y1,ddof=1)
varY2 = np.var(Y2,ddof=1)
varY3 = np.var(Y3,ddof=1)

print('computing covariances')
CovZY1 = np.sum([(z-meanZ)*(y-meanY1) for z,y in zip(Z,Y1)])/(N-1)
CovZY2 = np.sum([(z-meanZ)*(y-meanY2) for z,y in zip(Z,Y2)])/(N-1)
CovZY3 = np.sum([(z-meanZ)*(y-meanY3) for z,y in zip(Z,Y3)])/(N-1)

print('computing optimal alpha')
alpha1 = CovZY1/varY1
alpha2 = CovZY2/varY2
alpha3 = CovZY3/varY3

print('computing control variate samples')
Zcv1 = [z-alpha1*(y-meanY1) for z,y in zip(Z,Y1)]
Zcv2 = [z-alpha2*(y-meanY2) for z,y in zip(Z,Y2)]
Zcv3 = [z-alpha3*(y-meanY3) for z,y in zip(Z,Y3)]

print('Predicted reduction Y1 = ',(1 - CovZY1**2/varZ/varY1 ))
print('Actual reduction Y1 = ', np.var(Zcv1,ddof=1)/varZ)
print('Predicted reduction Y2 = ',(1 - CovZY2**2/varZ/varY2 ))
print('Actual reduction Y2 = ', np.var(Zcv2,ddof=1)/varZ)
print('Predicted reduction Y3 = ',(1 - CovZY3**2/varZ/varY3 ))
print('Actual reduction Y3 = ', np.var(Zcv3,ddof=1)/varZ)

############################################
#       Multiple Control Variate Case      #
############################################
# Compute covariance data 
CovYY = np.zeros((3,3))
for i in range(3):
    for j in range(i,3):
        meanYi = meanY[i]
        meanYj = meanY[j]
        CovYY[i,j] = np.sum([(yi-meanYi)*(yj-meanYj) for yi,yj in zip(Y[i],Y[j])])/(N-1)

for i in range(1,3):
    for j in range(i):
        CovYY[i,j] = CovYY[j,i]
CovZY = [CovZY1, CovZY2, CovZY3]

##### First Vector Control Variate
CovYY_sub = CovYY[:2,:2]
CovYY_sub_inv = np.linalg.inv(CovYY_sub)
CovZY_sub = CovZY[:2]
print('Predicted reduction (Y1, Y2) = ', 1-np.dot( CovZY_sub, np.dot(CovYY_sub_inv, CovZY_sub))/varZ )
alpha = np.dot( CovYY_sub_inv, CovZY_sub )
Zcv = [z-alpha[0]*(y1-meanY1)-alpha[1]*(y2-meanY2) for z,y1,y2 in zip(Z,Y1,Y2)]
print('Actual reduction (Y1, Y2) = ', np.var(Zcv)/varZ)

##### Second Vector Control Variate
CovYY_sub = CovYY[1:3,1:3]
CovYY_sub_inv = np.linalg.inv(CovYY_sub)
CovZY_sub = CovZY[1:3]
print('Predicted reduction (Y2, Y3) = ', 1-np.dot( CovZY_sub, np.dot(CovYY_sub_inv, CovZY_sub))/varZ )
alpha = np.dot( CovYY_sub_inv, CovZY_sub )
Zcv = [z-alpha[0]*(y2-meanY2)-alpha[1]*(y3-meanY3) for z,y2,y3 in zip(Z,Y2,Y3)]
print('Actual reduction (Y1, Y2) = ', np.var(Zcv)/varZ)

##### Third Vector Control Variate
CovYY_sub = CovYY[::2,::2]
CovYY_sub_inv = np.linalg.inv(CovYY_sub)
CovZY_sub = CovZY[::2]
print('Predicted reduction (Y1, Y3) = ', 1-np.dot( CovZY_sub, np.dot(CovYY_sub_inv, CovZY_sub))/varZ )
alpha = np.dot( CovYY_sub_inv, CovZY_sub )
Zcv = [z-alpha[0]*(y1-meanY1)-alpha[1]*(y3-meanY3) for z,y1,y3 in zip(Z,Y1,Y3)]
print('Actual reduction (Y1, Y2) = ', np.var(Zcv)/varZ)

##### Fourth Vector Control Variate
CovYY_inv = np.linalg.inv(CovYY)
print('Predicted reduction (Y1, Y2, Y3) = ', 1-np.dot( CovZY, np.dot(CovYY_inv, CovZY))/varZ )
alpha = np.dot( CovYY_inv, CovZY )
Zcv = [z-alpha[0]*(y1-meanY1)-alpha[1]*(y2-meanY2)-alpha[2]*(y3-meanY3) for z,y1,y2,y3 in zip(Z,Y1,Y2,Y3)]
print('Actual reduction (Y1, Y2) = ', np.var(Zcv)/varZ)