import pytest
import numpy as np

import matplotlib.pyplot as plt

from schroedinger import (
    Config, potential_interp, build_potential, solve_schroedinger
)


def psi(x, n, a):
    n += 1 # Index starting from 0
    x = -x + a / 2 # Reflect x and move to the left
    return np.sqrt(2 / a) * np.sin(n * np.pi * x / a)


def energy(n, a, m):
    return (n + 1)**2 * np.pi**2 / m / a**2 / 2.0


def test_infinite():
    conf = Config('test/infinite.inp')
    potential, delta = build_potential(conf)

    e, v = solve_schroedinger(conf.mass, potential[:, 1], delta, conf.eig_interval)
    # Account for -v also being an eigenvector if v is one
    v = np.abs(v)

    a = np.abs(conf.points[1][0] - conf.points[0][0])
    e_theory = np.zeros(e.shape)
    v_theory = np.zeros(v.shape)


    for n in range(e.shape[0]):
        e_theory[n] = energy(n, a, conf.mass)
        v_theory[:, n] = np.abs(psi(potential[:, 0], n, a))

    # for n in range(e.shape[0]):
    #     plt.plot(potential[:, 0], np.abs(v[:, n]), label='Num{}'.format(n))
    #     plt.plot(potential[:, 0], np.abs(v_theory[:, n]), ls='--', label='Theory{}'.format(n))
    # plt.legend()
    # plt.savefig('test.pdf')

    assert (np.allclose(e, e_theory, rtol=1e-2, atol=1e-2)
            and np.allclose(v, v_theory, rtol=1e-2, atol=1e-2))