import numpy as np
import scipy.stats as st

np.random.seed(42)

def brownian_bridge(t,a,b,Z):
    n=np.size(t)-2
    X=np.zeros(n+2)
    X[0]=a
    X[-1]=b
    tfinal=t[-1]
    #main loop
    for k in range(n):
    	mu = X[k] + (b-X[k])*(t[k+1]-t[k])/(tfinal-t[k])
    	sig2 = (tfinal-t[k])*(t[k+1]-t[k])/(tfinal-t[k])
    	X[k+1] = mu + np.sqrt(sig2)*Z[k]
    return X

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)

#gives some initial values 
X0 = 6
r = 0.05
sigma = 0.3
S = 10 
M = int(1e2) 
T = 1 
t = np.linspace(0,T,M)

##----------------------------------------------------------------------
#
#   Part 2.1 Monte Carlo
#
#----------------------------------------------------------------------
# pilot run
p = np.ones(S)/S #probabilties of each strata
si2 = np.zeros(S)
    
# pilot run for variance
Npilot = int(1e2)
Z = np.zeros(Npilot)
for i in range(Npilot):
    Z[i] = genPsi(st.norm.rvs(size=(M-1)),t,r,sigma,X0)
si2Z = np.var(Z)
print('Monte Carlo Var(Z): ',si2Z)

alpha = 1-0.99
tol = 1e-2
cval = st.norm.ppf(1-alpha/2,0,1)
NN = int(np.ceil( (cval*np.sqrt(si2Z)/tol)**2 ))
print(1,'tol = ',tol)
print(1,'\tN_MC: ',NN)
Z = np.zeros(NN)
for i in range(NN):
    Z[i] = genPsi(st.norm.rvs(size=(M-1)),t,r,sigma,X0)
print('\tMC estimator: ',  np.mean(Z))

##----------------------------------------------------------------------
#
#   Part 2.2  pilot run for optimal and proportional stratification
#
#-----------------------------------------------------------------------
print('\nPilot run for computing variances' )
print('-------------------------------------------------------------')
print('j  U_0 U_f    X_0             X_f  ')
# test run for the optimal allocation
muj = []
for j in range(S):
    X = np.zeros((Npilot,M))
    Nj = Npilot
    Zj = np.zeros(Nj)

    for i in range(Nj):
        U=(j+st.uniform.rvs())/S
        x_s=st.norm.ppf(U,0,np.sqrt(T))
        X[i,:]=brownian_bridge(t,0,x_s,st.norm.rvs(size=(M-1)))
        X[i,:]= X0*np.exp( (r-sigma**2/2)*t + sigma*X[i,:] )
        Zj[i] = np.max(X[i,:]) - np.min(X[i,:])
    si2[j] = np.var(Zj)
    print((j+1),j/S,(j+1)/S,
            st.norm.ppf(j/S,0,np.sqrt(T)),st.norm.ppf((j+1)/S,0,np.sqrt(T)))
    muj.append(np.mean(Zj))
print('Screening mean = ', np.dot(muj,p))
sumSiP = np.dot(p,np.sqrt(si2))
sumSi2P = np.dot(p,si2)
print('Var Opt = ', sumSiP**2)
print('Var Prop = ', sumSi2P)

##----------------------------------------------------------------------
#
#   Part 2.2.1  optimal stratification
#
#------------------------------------------------------------
print('\nOptimal stratification:' )
# Let's go
alpha = 1-0.99
tol = 1e-2
cval = st.norm.ppf(1-alpha/2,0,1)
NN = int(np.ceil( (cval*sumSiP/tol)**2 ))
print('tol = ',tol)
print('\tN_Str:',NN)
mu = np.zeros(S)
for j in range(S):
    X = np.zeros((NN,M))
    Nj = int(NN*p[j]*np.sqrt(si2[j])/sumSiP)
    Zj = np.zeros(Nj)
    for i in range(Nj):
        U=(j+st.uniform.rvs())/S
        x_s=st.norm.ppf(U,0,np.sqrt(T))
        X[i,:]=brownian_bridge(t,0,x_s,st.norm.rvs(size=(M-1)))
        X[i,:]= X0*np.exp( (r-sigma**2/2)*t + sigma*X[i,:] )
        Zj[i] = np.max(X[i,:]) - np.min(X[i,:])
    mu[j] = np.mean(Zj)
    print((j+1),Nj)


mu_str = np.dot(p,mu)
print('\tStrat. estimate: ',  mu_str)

##----------------------------------------------------------------------
#
#   Part 2.2.2  stratification  proportional
#
#------------------------------------------------------------

print('\nProportional stratification:' )
# Let's go
alpha = 1-0.99
tol = 1e-2
cval = st.norm.ppf(1-alpha/2,0,1)
NN = int(np.ceil( (cval/tol)**2*sumSi2P ))
print('tol = ',tol)
print('\tN_Str:',NN)
mu = np.zeros(S)
for j in range(S):
    X = np.zeros((NN,M))
    Nj = int(NN*p[j])
    Zj = np.zeros(Nj)
    for i in range(Nj):
        U=(j+st.uniform.rvs())/S
        x_s=st.norm.ppf(U,0,np.sqrt(T))
        X[i,:]=brownian_bridge(t,0,x_s,st.norm.rvs(size=(M-1)))
        X[i,:]= X0*np.exp( (r-sigma**2/2)*t + sigma*X[i,:] )
        Zj[i] = np.max(X[i,:]) - np.min(X[i,:])
    mu[j] = np.mean(Zj)
    print((j+1),Nj)
    
mu_str = np.dot(p,mu)
print('\tStrat. estimate: ',  mu_str)