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

def genPsi(xi,t,r,sigma,S0):
    dt=np.diff(t)
    W=np.concatenate([[0],np.cumsum(np.sqrt(dt)*xi)])
    S=S0*np.exp((r-0.5*sigma**2)*t+sigma*W)
    return np.max(S)-np.min(S)

T = 1.
M = 100
Nj = 2

sig = 0.3
r = 0.05
X0 = 6

ii=0
S_all = [10,50,100, 200, 500, 1000, 10000]
VS=np.zeros(len(S_all))
for S in S_all:
    pis = np.linspace(0, 1., S+1)[1:]
    #print pis
    strata = st.lognorm.ppf(pis,s=sig*np.sqrt(T), scale = np.exp(r-sig**2/2.))
    #print strata
    Y1 = r - sig**2/2. + sig * np.sqrt(T) * np.random.normal(size = (1000000,))
    X1 = np.exp(Y1)
    Xj = np.zeros((S, Nj))
    for i in range(S):
    	if i > 0:
            y = list(np.where((X1 > strata[i-1])*(X1 < strata[i]) == 1)[0])
    	else:
    	    y = list(np.where((X1 < strata[i]) == 1)[0])
    	Xj[i,:] = X1[y][:Nj]
    	Yj = np.log(Xj)
    Wj = (Yj - (r - sig**2/2.)) / sig
    Ypaths = np.zeros((M + 1, Wj.shape[0], Wj.shape[1]))
    dt = 1./(M)
    t = np.linspace(0, 1., M+1)
    for i in range(S):
    	for j in range(Nj):
            dw = np.sqrt(dt) * np.hstack([0,st.norm.rvs(size = (M,))])
            W = np.cumsum(dw)
            Ypaths[:,i,j] = (r - sig**2/2.)*t + sig * (W - t*(W[-1] - Wj[i,j]))
    Zj = X0*(np.exp(np.amax(Ypaths, axis=0)) - np.exp(np.amin(Ypaths, axis=0)))
    VS[ii]=np.var(Zj,axis = 1).sum() * (pis[1]-pis[0])**2 / 2.
    #print np.mean(Zj, axis = 1).sum() * (pis[1]-pis[0])
    print('NUMBER OF STRATA = ' + str(S))
    print('VARIANCE = ' + str(VS[ii]))
    ii+=1
    
plt.loglog(S_all,VS,label='Stratification')
plt.xlabel('Strata')
plt.ylabel('Variance')
plt.grid(True,which='both')

##----------------------------------------------------------------------
#
#    Monte Carlo
#
#----------------------------------------------------------------------
T = 1.
M = 100
Nj = 2
sig = 0.3
r = 0.05
X0 = 6
t = np.linspace(0, 1., M+1)

ii=0
S_all = [2*10,2*50,2*100, 2*200, 2*500, 2*1000, 2*10000]
VMC=np.zeros(len(S_all))
for NN in S_all:  
    # pilot run for variance
    Z=np.zeros(NN)
    for i in range(1,NN):
        Z[i-1] = genPsi(np.random.standard_normal(M),t,r,sig,X0);
    VMC[ii]= np.var(Z);
    ii+=1
       
plt.loglog(S_all,VMC/S_all,label='MC Variance')
plt.xlabel(r'$N=2S$')
plt.ylabel('Variance')
plt.grid(True,which='both')
plt.loglog(S_all,np.asarray(S_all)**-1.0,'k',linestyle='dashed',label=r'$N^{-1}$')
plt.legend()
plt.show()