import pytest
import numpy as np

import matplotlib.pyplot as plt

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

potential = {}
e = {}
v = {}

FORMS = ['asymmetric', 'double_linear', 'double_cubic', 'harmonic', 'finite']

@pytest.mark.parametrize('form', FORMS)
def test_potential(form: str) -> None:
    potential_prec = np.loadtxt(f'test/{form}/potential.dat')
    assert np.allclose(potential[form], potential_prec, rtol=1e-2, atol=1e-2)


@pytest.mark.parametrize('form', FORMS)
def test_e(form: str) -> None:
    e_prec = np.loadtxt(f'test/{form}/energies.dat')
    assert np.allclose(e[form], e_prec, rtol=1e-2, atol=1e-2)


@pytest.mark.parametrize('form', FORMS)
def test_v(form: str) -> None:
    v_prec = np.loadtxt(f'test/{form}/wavefuncs.dat')[:, 1:]
    assert np.allclose(v[form], v_prec, rtol=1e-2, atol=1e-2)


def setup_module():
    global conf
    global potential
    global e
    global v

    for form in FORMS:
        conf = Config(f'test/{form}.inp')
        potential[form], delta = build_potential(conf)
        e[form], v[form] = solve_schroedinger(conf.mass, potential[form][:, 1],
                                              delta, conf.eig_interval)