import numpy as np
import numpy.testing as npt

from numpy import (cos, exp, log, sin, sqrt)

from ..dopri5timeintegrator import Dopri5TimeIntegrator


class TestDopri5TimeIntegrator:
    def test__exponential_growth__works_ok(self):
        rhs = lambda y: y
        dt = 0.1
        tfinal = 1.0

        t = 0.0

        times = []
        solutions = []

        sut = Dopri5TimeIntegrator(rhs)
        sut.set_initial_value(1.0)

        t_step = 0
        while t < tfinal:
            t_step += 1
            t = t_step * dt

            sut.integrate(t)

            times.append(sut.t)
            solutions.append(sut.y[0])

        npt.assert_allclose(solutions, exp(times), rtol=1e-6)

    def test__exponential_decay__converges_to_3rd_order(self):
        tfinal = 5.0
        y0 = 1.0

        rhs_func = lambda y: -y

        powers = np.array(range(2, 11))
        numpoints = 2**powers
        errors = []

        for nums in numpoints:
            times, dt = np.linspace(0.0, tfinal, num=nums, retstep=True)
            times = times[1:]

            solutions = []

            # Use very coarse tolerance to measure convergence order.
            sut = Dopri5TimeIntegrator(rhs_func, rtol=1, atol=1)
            sut.set_initial_value(y0)

            for t in times:
                sut.integrate(t)
                solutions.append(sut.y[0])

            exact = exp(-times)
            errors.append(sqrt(dt) * np.linalg.norm(solutions - exact, 2))

        errors = np.array(errors)
        observed_orders = log(errors[0:-1]/errors[1:]) / log(2)

        expected_order = 4.0
        npt.assert_(np.all(observed_orders >= expected_order))

    def test__simple_oscillator__works_ok(self):
        k = 4.0
        m = 1.0
        omega = sqrt(k / m)

        y0, ydot0 = 1.0, 0.1
        ic = (y0, ydot0)

        tfinal = 1.0 + 0.09
        t = 0.0
        dt = 0.01

        rhs = lambda y: np.array([y[1], -omega**2 * y[0]])

        times = [t]
        solutions = [y0]

        sut = Dopri5TimeIntegrator(rhs)
        sut.set_initial_value(ic)

        t_step = 0
        while t < tfinal:
            t_step += 1
            t = t_step * dt

            sut.integrate(t)

            times.append(t)
            solutions.append(sut.y[0])

        times = np.array(times)

        c1, c2 = y0, ydot0/omega
        exact_solutions = c1 * cos(omega * times) + c2 * sin(omega * times)

        npt.assert_allclose(solutions, exact_solutions, rtol=1e-16, atol=1e-12)
