from numpy import exp, linspace, ones_like
from numpy.testing import assert_allclose

from ..characteristicstracer import CharacteristicsTracer


class TestCharacteristicsTracer:
    def test__velocity_is_constant__works_ok(self):
        RES = 32
        grid, dx = linspace(0, 1.0, num=RES+1, retstep=True)
        u = ones_like(grid)
        d = 0.0
        tfinal = 1.0
        dt = dx
        sut = CharacteristicsTracer(grid, u, d, tfinal, dt)

        x_actual, __ = sut.compute_characteristic(0.0)

        # Locations of the points on the characteristic curve must be
        # collocated with the grid points due to the chosen parameters.
        x_desired = grid.copy()

        assert_allclose(x_actual, x_desired, rtol=1e-16)

    def test__velocity_increases_linearly__trajectory_is_exponential(self):
        L = 5
        N = 200
        grid, dx = linspace(-L, 0.0, num=N+1, retstep=True)

        u = grid + 2*L
        d = 0.0

        tfinal = 1.0
        t, dt = linspace(0.0, tfinal, num=N+1, retstep=True)

        sut = CharacteristicsTracer(grid, u, d, tfinal, dt)

        x_actual, t_actual = sut.compute_characteristic(-L)
        x_desired = L * (-2 + exp(t_actual))

        # Somehow, I cannot get better match between actual and desired,
        # probably, due to the first-order Forward Euler integration in time.
        assert_allclose(x_actual, x_desired, rtol=1e-3, atol=2e-2)
