import pytest
from numericals import ode

import math

f1 = lambda x, y: y
f1_solution = lambda x: math.e**x
f1_params = (lambda x: math.e**x, 0, 1, 1, 1e-4)

f2 = lambda x, y: (6 * x) - (3 * y) + 5
f2_solution = lambda x: 2 * math.e ** (-3 * x) + 2 * x + 1
f2_params = (lambda x: 2 * math.e ** (-3 * x) + 2 * x + 1, 0, 5, 3, 1e-5)


def test_euler():
    f1_analytic, alpha, beta, y0, h = f1_params

    f1_euler = ode.euler(f1, alpha, beta, y0, h)
    f1_actual = [f1_analytic(i) for i in f1_euler[0]]

    assert f1_euler[1] == pytest.approx(f1_actual, abs=1e-3)

    f1_analytic, alpha, beta, y0, h = f2_params

    f2_euler = ode.euler(f2, alpha, beta, y0, h)
    f2_actual = [f1_analytic(i) for i in f2_euler[0]]

    assert f2_euler[1] == pytest.approx(f2_actual, abs=1e-3)


def test_heun():
    f1_analytic, alpha, beta, y0, h = f1_params

    f1_heun = ode.heun(f1, alpha, beta, y0, h)
    f1_actual = [f1_analytic(i) for i in f1_heun[0]]

    assert f1_heun[1] == pytest.approx(f1_actual, abs=1e-3)

    f2_analytic, alpha, beta, y0, h = f2_params

    f2_heun = ode.heun(f2, alpha, beta, y0, h)
    f2_actual = [f2_analytic(i) for i in f2_heun[0]]

    assert f2_heun[1] == pytest.approx(f2_actual, abs=1e-3)


def test_rk4():
    f1_analytic, alpha, beta, y0, h = f1_params

    f1_rk4 = ode.rk4(f1, alpha, beta, y0, h)
    f1_actual = [f1_analytic(i) for i in f1_rk4[0]]

    assert f1_rk4[1] == pytest.approx(f1_actual, abs=1e-3)

    f2_analytic, alpha, beta, y0, h = f2_params

    f2_rk4 = ode.rk4(f2, alpha, beta, y0, h)
    f2_actual = [f2_analytic(i) for i in f2_rk4[0]]

    assert f2_rk4[1] == pytest.approx(f2_actual, abs=1e-3)
