import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import root
from scipy.integrate import odeint, simps

fig, axs = plt.subplots(5, 1, figsize=(8, 30))

#question (b) and (c)
feven1=lambda z: np.tan(z)
feven2=lambda z: np.sqrt((z0/z)**2.0-1.0)
feven=lambda z: np.tan(z)-np.sqrt((z0/z)**2.0-1.0)
fodd1=lambda z: 1.0/np.tan(z)
fodd2=lambda z: -np.sqrt((z0/z)**2.0-1.0)
fodd=lambda z: 1.0/np.tan(z)+np.sqrt((z0/z)**2.0-1.0)
a=1.0
V0=20.0
m,hcut=1.0,1.0
z0=(a/2*hcut)*np.sqrt(2.0*m*V0)
z=np.linspace(1e-4,z0,1000,endpoint=False)

axs[0].plot(z,feven1(z),z,feven2(z))
axs[0].set_ylim(-10,10)
axs[0].grid()

#evenRootGuess=eval(input("Enter guess roots(z) for even parity (as list): "))
evenRootGuess=[1.2]

axs[1].plot(z,fodd1(z),z,fodd2(z))
axs[1].set_ylim(-10,10)
axs[1].grid()

#oddRootGuess=eval(input("Enter guess roots(z) for odd parity (as list): "))
oddRootGuess=[2.3]
evenRoots=root(feven,evenRootGuess).x
oddRoots=root(fodd,oddRootGuess).x

#question (d) and (e)
zRoots=np.sort(np.concatenate((evenRoots,oddRoots)))
eigenEnergies=(zRoots*hcut*2/a)**2/(2*m)-V0
print(eigenEnergies)
def V(x):
	return 0.0 if x<0 else -V0 if x<a else 0.0
def SE(psipsidot,x,E):
	psi,psidot=psipsidot
	psiddot=(2.0*m/hcut**2)*(V(x)-E)*psi
	return [psidot,psiddot]
def wavefunction(x,E):
	psi=odeint(SE,[0,1e-5],x,args=(E,))[:,0]
	psi=psi/np.sqrt(simps(psi**2,x))
	return psi

#question (f)
b1,b2=-2*a,3*a
x=np.linspace(b1,b2,1000)
def zeros(f):
	count=0
	for i in range(len(f)-1):
		if f[i]*f[i+1]<0:
			count=count+1
	return count
for energy in eigenEnergies:
	psi=wavefunction(x,energy)
	print(energy,zeros(psi),'odd' if zeros(psi)%2 else 'even')
	axs[2].plot(x,psi+energy,label=energy)
axs[2].plot(x,[V(i) for i in x]),
axs[2].legend()
axs[2].set_title('Plot of $\psi$ vs. x')

#question (g)
for energy in eigenEnergies:
	psi=wavefunction(x,energy)
	axs[3].plot(x,psi**2.0+energy,label='E={:.2f}'.format(energy))
axs[3].plot(x,[V(i) for i in x])
axs[3].set_title('Plot of $|\psi|^2$ vs. x')
axs[3].legend()

#question (i)
energy=eigenEnergies[0]
dE=energy*0.001/100
for i in range(3):
	psi=wavefunction(x,energy+(-1+i)*dE)
	axs[4].plot(x,psi)
axs[4].set_xlim(-3,5)
fig.savefig('fig.pdf')
fig.show()