import random
import numpy as np
import plotly.express as px
import matplotlib.pyplot as plt
import scipy as sc
import scipy.stats as stats

def simulate_1d_bm(nsteps=1000, t=0.001):
    steps = [ np.random.randn()*np.sqrt(t) for i in range(nsteps) ]
    steps[0] = 0
    y = np.cumsum(steps)
    x = [ t*i for i in range(nsteps) ]
    return x, y

#%%

nsims = 20
simulation_data = {}
z = np.linspace(0,1,1000)
fig = plt.figure()
fig.set_figheight(8)
fig.set_figwidth(12)
plt.plot(z, np.sqrt(2)*np.sqrt(z),'k'), plt.plot(z, -np.sqrt(2)*np.sqrt(z),'k')

finaltime = np.zeros(nsims)
firsttime = np.zeros(nsims)
secondtime = np.zeros(nsims)
thirdtime = np.zeros(nsims)
for i in range(nsims):
    x, y = simulate_1d_bm()
#    simulation_data['y{col}'.format(col=i)] = y
#    simulation_data['x'] = x
    finaltime[i] = y[-1]
    firsttime[i] = y[1]
    secondtime[i] = y[10]
    thirdtime[i] = y[20]
    plt.plot(x,y,'b')
    

plt.show()

#%%
#Histogram of values at $t=1$
plt.hist(finaltime, bins= np.int(np.sqrt(nsims)) + 1)
plt.show()

# Understandsing the limit
plt.plot(firsttime/np.sqrt(0.002)),
plt.plot(secondtime/np.sqrt(0.020)),
plt.plot(thirdtime/np.sqrt(0.040)), plt.show()
#Taking one solution and ploting the limit
x, y = simulate_1d_bm()
plt.plot(x,y/(np.sqrt(2)*np.sqrt(x))),plt.show

#%% Finding a good scale
fig = plt.figure()
fig.set_figheight(8)
fig.set_figwidth(12)


#Number of simulations
nsims=20
#Number of steps
steps = 100000
#Time normalized to 1
t = 1/steps

for i in range(nsims):
    x, y = simulate_1d_bm(steps,t)
#    simulation_data['y{col}'.format(col=i)] = y
#    simulation_data['x'] = x
    plt.plot(x,y,'b')

#We fix the scale of the x axis to be ten points 
plt.xlim(0, 10*t)
#The scale of the y axis is according by sqrt(t)
plt.ylim(-np.sqrt(10*t), np.sqrt(10*t) )
#we plot ten times more sqrt(2t) points to have a visual reference
z = np.linspace(0,1,10*steps)
plt.plot(z, np.sqrt(2)*np.sqrt(z),'k'), plt.plot(z, -np.sqrt(2)*np.sqrt(z),'k')
plt.show()



#%% Ej 15

q = 3
fig = plt.figure()
fig.set_figheight(8)
fig.set_figwidth(12)
for i in range(8):
#    x, y = simulate_1d_bm(2**(5+i), 1/2**(5+i))
    time_V = np.linspace(0,1,2**(5+i))
    
#    simulation_data['y{col}'.format(col=i)] = y
#    simulation_data['x'] = x
    
    V_steps = [np.abs(np.random.randn()*np.sqrt(1/2**(5+i)))**q for k in range(2**(5+i))]
    V = np.cumsum(V_steps)
    plt.plot(time_V,V)
    print(i,V[-1])

#%% Ej 16

#parte a)
def simulate_1d_bbridge(nsteps=1000, t=0.001):
    steps = [np.random.randn()*np.sqrt(t) for i in range(nsteps) ]
    steps[0] = 0
    aux = np.cumsum(steps)
    y= np.zeros_like(aux)
    for k in range(len(aux)):
        y[k] = aux[k] - k*t*aux[-1]
    x = [ t*i for i in range(nsteps) ]
    return x, y

fig = plt.figure()
fig.set_figheight(8)
fig.set_figwidth(12)
nsims = 10

# Parte b)

# First, we plot some trajectories with K> k_0, and the histogram for the K.
nsims = 1000
K = np.zeros(nsims)
k_0 = 1
cont = 0
for i in range(nsims):
    x, y = simulate_1d_bbridge()
    K[i] = np.max(np.abs(y))    
    if K[i]> k_0:
        plt.plot(x,y,'r')
        cont = cont + 1
    else:
        plt.plot(x,y,'k')

print(cont/nsims)
plt.show()

plt.hist(K, bins= np.int(np.sqrt(nsims)) + 1)
plt.show()

#%% Compute the expectation value (from several simulations), and the variance,

nsims = 20000
K = np.zeros(nsims)
cont = 0
for i in range(nsims):
    x, y = simulate_1d_bbridge()
    K[i] = np.max(np.abs(y))    



E_K = np.mean(K)
var_K = np.var(K) 

#%% plot the half-life distribution with a/b = E_K  and a/b^2 = var_K

theta = var_K/E_K
k = E_K/theta

x = np.linspace(0,3,100)
y = stats.gamma.pdf(x/theta, k, theta),
fig = plt.figure()
fig.set_figheight(8)
fig.set_figwidth(12)
plt.hist(K, bins= np.int(np.sqrt(nsims)) + 1, density='True')
plt.plot(x,y[0]/theta)
plt.show()

#%% Quantile 0.95:

print('k_0 s.t. P(K>k_0) =0.05: ',np.quantile(K,0.95))

print('Consistence with assume distribution: ',stats.gamma.cdf(np.quantile(K,0.95)/theta, k, theta))