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

N = 5
L = 1 # hva er L?????
V0_list = np.linspace(1.6E-19,5E-19,num = 10)
V0 = V0_list[9]
print(V0_list)
E_list = np.array([9.11213087913E-31,9.11172385727E-31,9.11144738959E-31,9.11124734388E-31,9.11109588069E-31,9.11097721845E-31,
                   9.11088174309E-31,9.11080326334E-31,9.110737612E-31,9.11068188132E-31])
E_list_right = (6.24150913E18)*np.sort(E_list)
T_list = np.array([0.775063189968,0.796204285259,0.824575490579,0.85665191579,0.889358375188,0.920129681519,
                   0.947395154217,0.968861578924,0.984525146436,0.992670126337])
T_analytic = np.zeros(len(V0_list))



#Potensialprofilen for en enkelt barriere
left = [0.0]*100*N
barrier = [V0]*N
right = [0.0]*100*N
V = np.asarray(left+barrier+right)
# Tabellen V har nå 201*N elementer


#Med Ntot punkter langs x-aksen, far vi Ntot ortonormerte egenfunksjoner
#og energiegenverdier ved  ̊a bruke numpyfunksjonen linalg.eigh,
#pa samme mate som i eksempelprogrammene som ligger tilgjengelig

#m=masse (SI)
hbar = 1.05E-34
m = 9.11E-31  #elektron
#dx = skrittlengde
dx = 1.0E-10

E0 = V0 + (np.pi)**2*V0/16
p0 = np.sqrt(2*m*E0)
k0 = p0/hbar
Ntot = len(V)

#distance = 101*N*dx
#middle_length = 10

#V = potensialet i boksen (V=0)
#d = liste med diagonalelementer i Hamiltonmatrisen H
d = [v + hbar**2/(m*dx**2) for v in V]
#e = verdi til ikke-diagonale elementer i H, dvs H(i,i+-1)
e = - hbar**2/(2*m*dx**2)



#Initialisering av matrisen H: Legger inn verdi 0 i samtlige elementer
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)

#Siden matrisen H er tridiagonal, finnes det helt sikkert mer e↵ektive funksjoner enn eigh.

#  En gaussformet starttilstand med senterposisjon midt i “venstre kontakt”
#og romlig utstrekning 1/10 av venstre kontakt
#kan for eksempel lages omtrent slik (se Oppgave 7 i øvingene):

#k0= 10 # k=skarpt definert bolgetall, endre denne verdien
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)
gaussinit = np.exp(-(x-x0)**2/(4*sigma**2))
planewavefactor = np.exp(1j*k0*x)
Psi0 = normfactor*gaussinit*planewavefactor


#Hvis for eksempel N = 10, har barrieren en tykkelse p ̊a 1.0 nm,
#mens de to kontaktene her har en bredde 100 nm.
#Planbølgefaktoren exp(ik0x) sørger for at bølgepakken har middelimpuls p0 = h ̄k0
#og middelhastighet v0 = p0/m = h ̄k0/m.





#Integral cj, regnes ut slik:
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)

#Forbereder figur, akser, og plotteelementet som skal animeres:
fig = plt.figure('Wave packet animation' , figsize=(16,8))
ymax = 10**8  #Ma justeres etter behov, eller beregnes
ax = plt.axes(xlim=(0, Ntot*dx), ylim=(0, ymax))
line, = ax.plot([], [], lw=1)
Vmax = np.argmax(V)   #riktig def???

#Initialisering; plotter bakgrunnen:
def init():
    line.set_data([], [])
    return line,


#Setter tidssteget, justeres etter behov, eller beregnes
tidssteg = 3.0E-15


def findRho():
    i = 13
    tidssteg = 3.0E-15
    t = i*tidssteg
    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
    return rho_t


mid_index = Ntot//2
E = hbar**2*(k0**2 + 1./(4.*sigma**2))/(2*E0)# Calculate energy expectation value [MeV]
rho_pluss = findRho()[mid_index:]
x_pluss = x[mid_index:]# Slice away list for x < 0
p_midlere = hbar*k0
I = np.trapz(rho_pluss,None,dx)          # Integrate with trapezoidal method
print("E: ", E)
print("T: ",I)

#Animasjonsfunksjon; kalles sekvensielt:
def animate(i):
    t = i*tidssteg
    #Beregner Psi(x,t)
    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
    line.set_data(x, rho_t)
    return line,

def T(V0,E):
    E0 = V0 + (np.pi)**2*V0/16
    #T = 1. / (1. + (V0**2/(4.*(V0-E)*E))*np.sinh(np.sqrt(2.*E0*(V0-E))/hbar*L)**2)
    T = (1 + (np.sinh(k0*L*np.sqrt(np.abs(E/V0-1)))**2)/(4*(E/V0)*(E/V0-1)))**(-1)
    return T


plt.plot(x,V*ymax/Vmax)
plt.xlabel('$x$ (m)',fontsize=20)
#Kaller animatoren.
#frames=antall bilder, dvs maxverdi for i;
#interval=varighet av hvert bilde i animasjonen m ̊alt i millisekunder
anim = animation.FuncAnimation(fig, animate, init_func=init, repeat=False,
                               frames=14, interval=1, blit=True)


'''for e in range(len(V0_list)):
    T_analytic[e] = T(V0_list[e],np.sort(E_list)[e])'''

fig2 = plt.figure("Transmission probability", figsize=(7,5),)
plt.plot(E_list_right,T_list, 'ro')
plt.xlabel('$E$ [eV]', fontsize=14)
plt.ylabel('$T(E)$', fontsize = 14)