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

iterations = 60000  # iterations for approximation
step = 0.0001  # step size for x
step_sqrd = pow(step, 2)  # square of step size
epsilon = 2.5  # 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

fig, ax = plt.subplots()
line, = ax.plot([], [], label=f'epsilon = {epsilon}')

ax.set_xlim(-1, 1)
ax.set_ylim(-1, 1)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_title("Schrodinger Eqn in Harmonic Potential")
ax.legend(loc=1)


def init():
    line.set_data([], [])
    return line,


def update(frame):
    global pos, potential, psi, potential_past_2, potential_past_1, psi_past_2, psi_past_1
    for i in range(frame * 1000, (frame * 1000) + 1000):
        print(i)
        pos += step
        potential = 2 * epsilon - pow(pos, 2)  # 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

    line.set_data(x_out, y_out)
    return line,


ani = animation.FuncAnimation(fig, update, frames=np.arange(0, (iterations - 2)), init_func=init, blit=True)
plt.show()