import pylab as lab
import numpy as np

iterations = 60000  # iterations for approximation

step = 0.0001  # step size for x
step_sqrd = pow(step, 2)  # square of step size

well_width = 2

# try values 1, 2, and 3 for various solutions
n = 1
epsilon = (n * np.pi) ** 2 / (well_width ** 2)  # energy level, should be integer n+1/2 for good solutions

psi = 0.0  # initial value of wave function
potential = 0.0  # initial value of the potential energy function
pos = -1 * (iterations - 2) * step  # initial value of the position

potential_past_2 = epsilon + pos - 2 * step  # k_0, potential energy at two steps before current
potential_past_1 = epsilon + pos - step  # k_1, potential energy at one step before current
amplitude = 0.1  # initial amplitude of wave function
psi_past_2 = 0  # y_0, wave function at two steps before current
psi_past_1 = amplitude  # y_1, wave function at one step before current

x_out = []  # list to store x values for plotting
y_out = []  # list to store y values for plotting

count = -1 * iterations + 2  # counter for the loop

# Numerov integration loop
while count < iterations - 2:
    count += 1
    pos += step
    # close enough to infinity :)
    potential = 100000000
    if abs(pos) < well_width / 2:
        potential = epsilon  # potential energy function
    b = step_sqrd / 12  # constant used for Numerov
    # Numerov method to calculate psi at current step
    psi = ((2 * (1 - 5 * b * potential_past_1) * psi_past_1 - (1 + b * potential_past_2) * psi_past_2)
           / (1 + b * potential))

    # Save for plotting
    x_out.append(pos)
    y_out.append(psi)

    # Shift for next iteration
    psi_past_2 = psi_past_1
    psi_past_1 = psi
    potential_past_2 = potential_past_1
    potential_past_1 = potential

# Plot
lab.figure(1)
lab.plot(x_out, y_out, label=f'epsilon = {epsilon}')
lab.xlabel("x")
lab.ylabel("y")
lab.title("Schrodinger Eqn in Harmonic Potential")
lab.legend(loc=1)
lab.show()