from __future__ import division, print_function, absolute_import
from itertools import product
from numpy.testing import (assert_, assert_allclose,
                           assert_equal, assert_no_warnings)
import pytest
from pytest import raises as assert_raises
from scipy._lib._numpy_compat import suppress_warnings
import numpy as np
from scipy.optimize._numdiff import group_columns
from scipy.integrate import solve_ivp, RK23, RK45, Radau, BDF, LSODA
from scipy.integrate import OdeSolution
from scipy.integrate._ivp.common import num_jac
from scipy.integrate._ivp.base import ConstantDenseOutput
from scipy.sparse import coo_matrix, csc_matrix


def fun_linear(t, y):
    return np.array([-y[0] - 5 * y[1], y[0] + y[1]])


def jac_linear():
    return np.array([[-1, -5], [1, 1]])


def sol_linear(t):
    return np.vstack((-5 * np.sin(2 * t),
                      2 * np.cos(2 * t) + np.sin(2 * t)))


def fun_rational(t, y):
    return np.array([y[1] / t,
                     y[1] * (y[0] + 2 * y[1] - 1) / (t * (y[0] - 1))])


def fun_rational_vectorized(t, y):
    return np.vstack((y[1] / t,
                      y[1] * (y[0] + 2 * y[1] - 1) / (t * (y[0] - 1))))


def jac_rational(t, y):
    return np.array([
        [0, 1 / t],
        [-2 * y[1] ** 2 / (t * (y[0] - 1) ** 2),
         (y[0] + 4 * y[1] - 1) / (t * (y[0] - 1))]
    ])


def jac_rational_sparse(t, y):
    return csc_matrix([
        [0, 1 / t],
        [-2 * y[1] ** 2 / (t * (y[0] - 1) ** 2),
         (y[0] + 4 * y[1] - 1) / (t * (y[0] - 1))]
    ])


def sol_rational(t):
    return np.asarray((t / (t + 10), 10 * t / (t + 10) ** 2))


def fun_medazko(t, y):
    n = y.shape[0] // 2
    k = 100
    c = 4

    phi = 2 if t <= 5 else 0
    y = np.hstack((phi, 0, y, y[-2]))

    d = 1 / n
    j = np.arange(n) + 1
    alpha = 2 * (j * d - 1) ** 3 / c ** 2
    beta = (j * d - 1) ** 4 / c ** 2

    j_2_p1 = 2 * j + 2
    j_2_m3 = 2 * j - 2
    j_2_m1 = 2 * j
    j_2 = 2 * j + 1

    f = np.empty(2 * n)
    f[::2] = (alpha * (y[j_2_p1] - y[j_2_m3]) / (2 * d) +
              beta * (y[j_2_m3] - 2 * y[j_2_m1] + y[j_2_p1]) / d ** 2 -
              k * y[j_2_m1] * y[j_2])
    f[1::2] = -k * y[j_2] * y[j_2_m1]

    return f


def medazko_sparsity(n):
    cols = []
    rows = []

    i = np.arange(n) * 2

    cols.append(i[1:])
    rows.append(i[1:] - 2)

    cols.append(i)
    rows.append(i)

    cols.append(i)
    rows.append(i + 1)

    cols.append(i[:-1])
    rows.append(i[:-1] + 2)

    i = np.arange(n) * 2 + 1

    cols.append(i)
    rows.append(i)

    cols.append(i)
    rows.append(i - 1)

    cols = np.hstack(cols)
    rows = np.hstack(rows)

    return coo_matrix((np.ones_like(cols), (cols, rows)))


def fun_complex(t, y):
    return -y


def jac_complex(t, y):
    return -np.eye(y.shape[0])


def jac_complex_sparse(t, y):
    return csc_matrix(jac_complex(t, y))


def sol_complex(t):
    y = (0.5 + 1j) * np.exp(-t)
    return y.reshape((1, -1))


def compute_error(y, y_true, rtol, atol):
    e = (y - y_true) / (atol + rtol * np.abs(y_true))
    return np.sqrt(np.sum(np.real(e * e.conj()), axis=0) / e.shape[0])


def test_integration():
    rtol = 1e-3
    atol = 1e-6
    y0 = [1/3, 2/9]

    for vectorized, method, t_span, jac in product(
            [False, True],
            ['RK23', 'RK45', 'Radau', 'BDF', 'LSODA'],
            [[5, 9], [5, 1]],
            [None, jac_rational, jac_rational_sparse]):

        if vectorized:
            fun = fun_rational_vectorized
        else:
            fun = fun_rational

        with suppress_warnings() as sup:
            sup.filter(UserWarning,
                       "The following arguments have no effect for a chosen solver: `jac`")
            res = solve_ivp(fun, t_span, y0, rtol=rtol,
                            atol=atol, method=method, dense_output=True,
                            jac=jac, vectorized=vectorized)
        assert_equal(res.t[0], t_span[0])
        assert_(res.t_events is None)
        assert_(res.success)
        assert_equal(res.status, 0)

        assert_(res.nfev < 40)

        if method in ['RK23', 'RK45', 'LSODA']:
            assert_equal(res.njev, 0)
            assert_equal(res.nlu, 0)
        else:
            assert_(0 < res.njev < 3)
            assert_(0 < res.nlu < 10)

        y_true = sol_rational(res.t)
        e = compute_error(res.y, y_true, rtol, atol)
        assert_(np.all(e < 5))

        tc = np.linspace(*t_span)
        yc_true = sol_rational(tc)
        yc = res.sol(tc)

        e = compute_error(yc, yc_true, rtol, atol)
        assert_(np.all(e < 5))

        tc = (t_span[0] + t_span[-1]) / 2
        yc_true = sol_rational(tc)
        yc = res.sol(tc)

        e = compute_error(yc, yc_true, rtol, atol)
        assert_(np.all(e < 5))

        # LSODA for some reasons doesn't pass the polynomial through the
        # previous points exactly after the order change. It might be some
        # bug in LSOSA implementation or maybe we missing something.
        if method != 'LSODA':
            assert_allclose(res.sol(res.t), res.y, rtol=1e-15, atol=1e-15)


def test_integration_complex():
    rtol = 1e-3
    atol = 1e-6
    y0 = [0.5 + 1j]
    t_span = [0, 1]
    tc = np.linspace(t_span[0], t_span[1])
    for method, jac in product(['RK23', 'RK45', 'BDF'],
                               [None, jac_complex, jac_complex_sparse]):
        with suppress_warnings() as sup:
            sup.filter(UserWarning,
                       "The following arguments have no effect for a chosen solver: `jac`")
            res = solve_ivp(fun_complex, t_span, y0, method=method,
                            dense_output=True, rtol=rtol, atol=atol, jac=jac)

        assert_equal(res.t[0], t_span[0])
        assert_(res.t_events is None)
        assert_(res.success)
        assert_equal(res.status, 0)

        assert_(res.nfev < 25)
        if method == 'BDF':
            assert_equal(res.njev, 1)
            assert_(res.nlu < 6)
        else:
            assert_equal(res.njev, 0)
            assert_equal(res.nlu, 0)

        y_true = sol_complex(res.t)
        e = compute_error(res.y, y_true, rtol, atol)
        assert_(np.all(e < 5))

        yc_true = sol_complex(tc)
        yc = res.sol(tc)
        e = compute_error(yc, yc_true, rtol, atol)

        assert_(np.all(e < 5))


def test_integration_sparse_difference():
    n = 200
    t_span = [0, 20]
    y0 = np.zeros(2 * n)
    y0[1::2] = 1
    sparsity = medazko_sparsity(n)

    for method in ['BDF', 'Radau']:
        res = solve_ivp(fun_medazko, t_span, y0, method=method,
                        jac_sparsity=sparsity)

        assert_equal(res.t[0], t_span[0])
        assert_(res.t_events is None)
        assert_(res.success)
        assert_equal(res.status, 0)

        assert_allclose(res.y[78, -1], 0.233994e-3, rtol=1e-2)
        assert_allclose(res.y[79, -1], 0, atol=1e-3)
        assert_allclose(res.y[148, -1], 0.359561e-3, rtol=1e-2)
        assert_allclose(res.y[149, -1], 0, atol=1e-3)
        assert_allclose(res.y[198, -1], 0.117374129e-3, rtol=1e-2)
        assert_allclose(res.y[199, -1], 0.6190807e-5, atol=1e-3)
        assert_allclose(res.y[238, -1], 0, atol=1e-3)
        assert_allclose(res.y[239, -1], 0.9999997, rtol=1e-2)


def test_integration_const_jac():
    rtol = 1e-3
    atol = 1e-6
    y0 = [0, 2]
    t_span = [0, 2]
    J = jac_linear()
    J_sparse = csc_matrix(J)

    for method, jac in product(['Radau', 'BDF'], [J, J_sparse]):
        res = solve_ivp(fun_linear, t_span, y0, rtol=rtol, atol=atol,
                        method=method, dense_output=True, jac=jac)
        assert_equal(res.t[0], t_span[0])
        assert_(res.t_events is None)
        assert_(res.success)
        assert_equal(res.status, 0)

        assert_(res.nfev < 100)
        assert_equal(res.njev, 0)
        assert_(0 < res.nlu < 15)

        y_true = sol_linear(res.t)
        e = compute_error(res.y, y_true, rtol, atol)
        assert_(np.all(e < 10))

        tc = np.linspace(*t_span)
        yc_true = sol_linear(tc)
        yc = res.sol(tc)

        e = compute_error(yc, yc_true, rtol, atol)
        assert_(np.all(e < 15))

        assert_allclose(res.sol(res.t), res.y, rtol=1e-14, atol=1e-14)


@pytest.mark.slow
@pytest.mark.parametrize('method', ['Radau', 'BDF', 'LSODA'])
def test_integration_stiff(method):
    rtol = 1e-6
    atol = 1e-6
    y0 = [1e4, 0, 0]
    tspan = [0, 1e8]

    def fun_robertson(t, state):
        x, y, z = state
        return [
            -0.04 * x + 1e4 * y * z,
            0.04 * x - 1e4 * y * z - 3e7 * y * y,
            3e7 * y * y,
        ]

    res = solve_ivp(fun_robertson, tspan, y0, rtol=rtol,
                    atol=atol, method=method)

    # If the stiff mode is not activated correctly, these numbers will be much bigger
    assert res.nfev < 5000
    assert res.njev < 200


def test_events():
    def event_rational_1(t, y):
        return y[0] - y[1] ** 0.7

    def event_rational_2(t, y):
        return y[1] ** 0.6 - y[0]

    def event_rational_3(t, y):
        return t - 7.4

    event_rational_3.terminal = True

    for method in ['RK23', 'RK45', 'Radau', 'BDF', 'LSODA']:
        res = solve_ivp(fun_rational, [5, 8], [1/3, 2/9], method=method,
                        events=(event_rational_1, event_rational_2))
        assert_equal(res.status, 0)
        assert_equal(res.t_events[0].size, 1)
        assert_equal(res.t_events[1].size, 1)
        assert_(5.3 < res.t_events[0][0] < 5.7)
        assert_(7.3 < res.t_events[1][0] < 7.7)

        event_rational_1.direction = 1