import numpy as np
from matplotlib import pyplot as plt
import matplotlib.animation as animation

hbar=1.054718E-34
N=10
V0=2.0E-19
m=9.11E-31
dx = 1.0E-10
Vmax=13.0E-20
tidssteg=5.0E-15
iterations=5
b=15
p0 = 3.2*np.pi*hbar/(b*dx) #resonant energi
#p0 = np.pi*hbar*2.9/(b*dx) #ikke-resonant energi
k0=p0/hbar

#Område der vi skal finne P(t)
vens = 101*N
hoyr = (101+b)*N
vlen = vens*dx
hlen = hoyr*dx

#definerer potensialene
left=[0.0]*100*N
barrier_left=[V0]*N
between=[0.0]*b*N
barrier_right=[V0]*N
right=[0.0]*100*N
V=np.asarray(left+barrier_left+between+barrier_right+right)

#stasjonære tilstander
d = [v + hbar**2/(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)


#initialtilstand
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)
normfactor = (2*np.pi*sigma**2)**(-0.25)
gaussinitial = np.exp(-(x-x0)**2/(4*sigma**2))
planewavefactor = np.exp(1j*k0*x)
psi0 = normfactor*gaussinitial*planewavefactor


#bølgepakke som lineær komb. av stasj. tilstander
psi_matrix_complex = psi_matrix*(1.0 + 0.0j)
c=np.zeros(Ntot, dtype = np.complex128)
for n in range(Ntot):
    c[n]=np.vdot(psi_matrix_complex[:,n],psi0)

#finner andel av bølgen som befinner seg mellom veggene
probabilitylist = []
frames = 250
tlist = [n*tidssteg for n in range(frames)]
for i in range(frames):
    t = i*tidssteg
    psi_t=np.zeros(Ntot, dtype=np.complex128)
    rhosum_enclosed = 0
    for n in range(Ntot):
        psi_t=psi_t+c[n]*psi_matrix_complex[:,n]*np.exp(-1j*energy[n]*t/hbar)
        if n>=vens and n<=hoyr:
            rhosum_enclosed += np.abs(psi_t[n])**2 #mengde mellom vegger
    rho_t=np.abs(psi_t)**2
    total_rho = np.sum(rho_t) #hele rho_t
    probabilitylist.append(rhosum_enclosed/total_rho)

def trapezoid(f, xlist):
    Sum = 0
    for i in range(len(xlist)-1):
        Sum += (dx/2)*(f[i]+f[i+1])
    return Sum

"""
plist=np.linspace(p0/4, p0*1.1, 60)

Tlist=[]
energylist=[]

for i in range(60):
  p_temp = plist[i]
  k0=p_temp/hbar
  t=(tidssteg)

  planewavefactor = np.exp(1j*k0*x)
  psi0 = normfactor*gaussinitial*planewavefactor
  c=np.zeros(Ntot, dtype = np.complex128)
  for n in range(Ntot):
      c[n]=np.vdot(psi_matrix_complex[:,n],psi0)
  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]*t/hbar)

  rho_t=np.abs(psi_t)**2
  R = trapezoid(rho_t, left)
  T=1-R
  Tlist.append(T)
  energylist.append(p_temp*p_temp/(2*m))

plt.figure()

plt.xlabel('E', fontsize=16)
plt.ylabel('T(E)', fontsize=16)
plt.plot(energylist, Tlist)
plt.savefig('Transsans(E).pdf')
plt.show()
"""

plt.figure('P(t)', figsize=(20,20))
plt.title('Sannsynlighet for å finne elektron mellom barrierene, resonans', fontsize=22)
#plt.title('Sannsynlighet for å finne elektron mellom barrierene, ikke resonans', fontsize=22)
plt.plot(tlist, probabilitylist)
plt.ylabel('P_b(t)',fontsize=20)
plt.xlabel('t', fontsize=20)
plt.savefig('Doublebarrier_resonance.pdf')
#plt.savefig('Doublebarrier_nonresonance.pdf')
plt.show()