# -*- coding: utf-8 -*-
"""
Spyder Editor

This is a temporary script file.
"""
from scipy import *
import numpy as np
import matplotlib.pyplot as plt

#Task 1

xmin=-8    #-16 for excited state
xmax=-xmin
Nsteps=1001 #2001 for excited state
alpha=1.0
VO=4.0

X=np.linspace(xmin,xmax,Nsteps)
Y=np.linspace(0,0,Nsteps)
h=(X[1]-X[0])

B=[]

for i in range(Nsteps):
    B.append((-VO)*np.exp(-alpha*(X[i]**2)))
Vx=np.asarray(B)
E=-2.774534188


k=np.sqrt(-2*E)
Y[0]=1
Y[1]=np.exp(k*h)
C=[]
Fx=2*(Vx-E)



for i in range(Nsteps-2):
    Y[i+2]=(2+(h**2)*Fx[i+1])*Y[i+1]-Y[i]
plt.plot(X,Y)

"""
##Task 2 a)
xO=1.5
xmin=-16
xmax=-xmin
Nsteps=2001
alpha=1.0
VO=2.0

X=np.linspace(xmin,xmax,Nsteps)
Y=np.linspace(0,0,Nsteps)
h=(X[1]-X[0])

B=[]

for i in range(Nsteps):
    B.append((-VO)*((np.exp(-alpha*((X[i]-xO)**2)))+(np.exp(-alpha*((X[i]+xO)**2)))))
Vx=np.asarray(B)
E=-1.259666848796
k=np.sqrt(-2*E)
Y[0]=1
Y[1]=np.exp(k*h)
C=[]
Fx=2*(Vx-E)



for i in range(Nsteps-2):
    Y[i+2]=(2+(h**2)*Fx[i+1])*Y[i+1]-Y[i]
plt.plot(X,Y)


#Task 2 b)
xO=1.5
xmin=-16
xmax=-xmin
Nsteps=2001
alpha=1.0
VO=4.0

X=np.linspace(xmin,xmax,Nsteps)
Y=np.linspace(0,0,Nsteps)
h=(X[1]-X[0])

B=[]

for i in range(Nsteps):
    B.append((-VO)*((np.exp(-alpha*((X[i]-xO)**2)))+(np.exp(-alpha*((X[i]+xO)**2)))))
Vx=np.asarray(B)
E=-0.586826045
k=np.sqrt(-2*E)
Y[0]=1
Y[1]=np.exp(k*h)
C=[]
Fx=2*(Vx-E)



for i in range(Nsteps-2):
    Y[i+2]=(2+(h**2)*Fx[i+1])*Y[i+1]-Y[i]
plt.plot(X,Y)

#A and W
E_g = [-3.101693348, -1.251364587, -2.890853493, -1.130936465, -2.813694387, -1.005513342]
E_e = [-2.715730925, -0.376767602, -2.736275220, -0.49283514, -2.754089237, -0.586826045]
E0_e=-0.78401783878
E0_g=-2.774534188
W=[]
A=[]
for i in range(6):
    if i%2 ==0:
        W.append((E_g[i]-E_e[i])/2)
        A.append(E0_g - E_g[i] + W[i])
    else:
        W.append((E_g[i]-E_e[i])/2)
        A.append(E0_e - E_g[i] + W[i])

print(W)
print(A)
"""