#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm
from matplotlib import rc
from numpy import matlib
####################################################################
rc('font',**{'family':'serif','serif':['Computer Modern Roman'],
     'size' : '12'})
rc('text', usetex=True)
rc('lines', linewidth=2)
plt.rcParams['axes.facecolor']='w'
import matplotlib
latex_preamble = r'\usepackage{amsmath} \usepackage{amssymb}'
matplotlib.rcParams.update({
    'text.usetex': True,
    'text.latex.preamble': latex_preamble
})

####################################################################
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+1])*(t[k+1]-t[k])/(tfinal-t[k])
        X[k+1] = mu +np.sqrt(sig2)*Z[k]
    return X

S = 12 # number of strata
N = [2]*S # list of no of samples per stratum
M = int(1e3) # number of points for path construction
T = 1 # terminal time at which we stratify
t = np.linspace(0,T,M)
print('j\tU_lo\tU_up\tX_lo\tX_up')

for j in range(S):
    print('%d'% (j+1),'\t%.2f' % (j/S),'\t%.2f' % ((j+1)/S),'\t%.2f' % norm.ppf(j/S) ,'\t%.2f' % norm.ppf((j+1)/S))
    plt.plot([0,1],[norm.ppf(j/S,0,np.sqrt(T)), norm.ppf(j/S,0,np.sqrt(T))],'k:')
    Nj = N[j]
    for i in range(Nj):
        U=(j+np.random.random(1))/S # sample uniform random number from Strata j
        x_s=norm.ppf(U,0,np.sqrt(T)) # invert it to get the stratified terminal value
        B=brownian_bridge(t,0,x_s,np.random.standard_normal(M))
      
        plt.plot(t,B,'C'+str(j))

plt.show()