#plotting probability for finding a particle between barriers
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from IPython.display import HTML


N = 10
V0 = 2e-19
hbar = 1.05e-34
m = 9.11e-31
dx = 1e-10
left = [0.0]*100*N
barrier = [V0]*N
barrier2 = [V0]*N
middle = [0.0]*50*N
right = [0.0]*100*N


V = np.asarray(left+barrier+middle+barrier2+right)

#making psi_n matrix and energy eigenvalues solving TISE
d = [v + hbar**2/(m*dx**2) for v in V]
e = - hbar**2/(2*m*dx**2)
Ntot=len(V)
H = [[0]*(Ntot) for n in range(Ntot)]
for i in range(Ntot):
    for j in range(Ntot):
        if i==j:
            H[i][j]=d[i]
        if abs(i-j)==1:
            H[i][j]=e
H = np.asarray(H)
energy,psi_matrix = np.linalg.eigh(H)

def wave_maker(E):
    m = 9.11e-31
    hbar = 1.05e-34
    dx = 1e-10
    x = np.asarray([dx*n for n in range(Ntot)])
    x0 = 50*N*dx
    sigma = 10*N*dx
    Delta_x = sigma
    Delta_p = hbar/(2.0*sigma)
    #k0 = p0/hbar
    k0 = np.sqrt(2*m*E)/hbar
    normfactor = (2*np.pi*sigma**2)**(-0.25)
    gaussinit = np.exp(-(x-x0)**2/(4*sigma**2))
    psi_matrix_complex = psi_matrix*(1.0 + 0.0j)

    planewavefactor = np.exp(1j*k0*x)
    Psi0 = normfactor*gaussinit*planewavefactor
    c = np.zeros(Ntot, dtype = np.complex128)
    for n in range(Ntot):
        c[n] = np.vdot(psi_matrix_complex[:,n],Psi0)

    t = 240*m*N*dx/(k0*hbar)
    timelist = np.linspace(t/20,t*20,200)
    problist = [0]*len(timelist)
    for i in range(len(timelist)):
        Psi_t = np.zeros(Ntot,dtype=np.complex128)
        for n in range(Ntot):
            Psi_t = Psi_t + c[n]*psi_matrix_complex[:,n]*np.exp(-1j*energy[n]*timelist[i]/hbar)
        rho_t = np.abs(Psi_t)**2
        M = np.sum(rho_t[len(left):len(middle)+len(left)]*dx)/sum(rho_t*dx)
        problist[i]=M
    return timelist,problist


p0 = 1e-24
E = p0**2/(2*m)

timelist,problist = wave_maker(E)

plt.figure()
plt.title("Sannsynlighet P for å finne partikkel mellom to barrierer som funksjon av tida, uten resonans")
plt.plot(timelist,problist, label = "P(t)")
plt.show()