from qutip.solver.dysolve_propagator import DysolvePropagator, dysolve_propagator
from qutip.solver import propagator
from qutip.solver.cy.dysolve import cy_compute_integrals
from qutip import (
    sigmax, sigmay, sigmaz, qeye, qeye_like, tensor, enr_destroy, CoreOptions
)
from scipy.special import factorial
import numpy as np
import pytest


@pytest.fixture(scope='module')
def empty_instance():
    return DysolvePropagator.__new__(DysolvePropagator)


def _enr_xx():
    a, b = enr_destroy([2, 2], 1)
    return (a + a.dag()) @ (b + b.dag())

def _enr_xz():
    a, b = enr_destroy([2, 2], 1)
    return (a + a.dag()) @ (b.dag() @ b)

def _enr_zz():
    a, b = enr_destroy([2, 2], 1)
    return (a.dag() @ a) @ (b.dag() @ b)


@pytest.mark.parametrize("eff_omega", [-10.0, -1.0, -0.1, 0.1, 1.0, 10.0])
@pytest.mark.parametrize("dt", [-10.0, -1.0, -0.1, 0.1, 1.0, 10.0])
@pytest.mark.parametrize("ws, answer", [
    # First part of tuple is "ws", second part is "answer"
    (
        np.array([0.0]),
        lambda _, dt: dt
    ),
    (
        np.array([1e-12]),
        lambda _, dt: dt
    ),
    (
        lambda eff_omega: np.array([eff_omega]),
        lambda eff_omega, dt: (-1j/eff_omega) * (np.exp(1j*eff_omega*dt) - 1)
    ),
    (
        np.array([0.0, 0.0, 0.0, 0.0, 0.0]),
        lambda _, dt: (dt**5) / factorial(5)
    ),
    (
        np.array([1e-12, 1e-12, 1e-12]),
        lambda _, dt: (dt**3) / factorial(3)
    ),
    (
        lambda eff_omega: np.array([eff_omega, 0.0]),
        lambda eff_omega, dt: (-1j/eff_omega) * (
            (-1j/eff_omega) * (np.exp(1j*eff_omega*dt) - 1) - dt
        )
    ),
    (
        lambda eff_omega: np.array([0.0, eff_omega]),
        lambda eff_omega, dt: (-1j*dt/eff_omega) * np.exp(1j*eff_omega*dt) -
        ((1j/eff_omega)**2) * (np.exp(1j*eff_omega*dt)-1))
])
def test_integrals_1(eff_omega, dt, ws, answer):
    if callable(ws):
        ws = ws(eff_omega)
    if callable(answer):
        answer = answer(eff_omega, dt)

    integrals = cy_compute_integrals(ws, dt)

    assert np.isclose(integrals, answer, rtol=1e-10, atol=1e-10)


@pytest.mark.parametrize("eff_omega_1", [-25.0, -5.0, -0.5, 0.5, 5.0, 25.0])
@pytest.mark.parametrize("eff_omega_2", [-25.0, -5.0, -0.5, 0.5, 5.0, 25.0])
@pytest.mark.parametrize("dt", [-10.0, -1.0, -0.1, 0.1, 1.0, 10.0])
def test_integrals_2(eff_omega_1, eff_omega_2, dt):
    ws = np.array([eff_omega_1, eff_omega_2])
    integrals = cy_compute_integrals(ws, dt)

    if eff_omega_1 + eff_omega_2 == 0:
        answer = (-1j*dt/eff_omega_1) + \
            (np.exp(1j*eff_omega_2*dt)-1)/(eff_omega_1*eff_omega_2)
    else:
        exp_1 = np.exp(1j*(eff_omega_1+eff_omega_2)*dt)
        exp_2 = np.exp(1j*eff_omega_2*dt)
        answer = -(exp_1-1)/(eff_omega_1*(eff_omega_1+eff_omega_2)) + \
            (exp_2-1)/(eff_omega_1*eff_omega_2)

    assert np.isclose(integrals, answer, rtol=1e-10, atol=1e-10)


@pytest.mark.parametrize("max_order, X, answer", [
    (
        1,
        sigmay(),
        np.array([0, -1j, 1j, 0])
    ),
    (
        2,
        sigmay(),
        np.array([0, 0, 1, 0, 0, 1, 0, 0])
    ),
    (
        3,
        sigmay(),
        np.array([0, 0, 0, 0, 0, -1j, 0, 0, 0, 0, 1j, 0, 0, 0, 0, 0])
    ),
    (
        1,
        sigmax(),
        np.array([0, 1, 1, 0])
    ),
    (
        2,
        sigmax(),
        np.array([0, 0, 1, 0, 0, 1, 0, 0])
    ),
    (
        3,
        sigmax(),
        np.array([0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0])
    ),
    (
        1,
        sigmaz(),
        np.array([1, 0, 0, -1])
    ),
    (
        2,
        sigmaz(),
        np.array([1, 0, 0, 0, 0, 0, 0, 1])
    ),
    (
        3,
        sigmaz(),
        np.array([1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -1])
    ),
    (
        1,
        qeye(2),
        np.array([1, 0, 0, 1])
    ),
    (
        2,
        qeye(2),
        np.array([1, 0, 0, 0, 0, 0, 0, 1])
    ),
    (
        3,
        qeye(2),
        np.array([1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1])
    ),
    (
        1,
        qeye(3),
        np.array([1, 0, 0, 0, 1, 0, 0, 0, 1])
    ),
    (
        2,
        qeye(3),
        np.array([1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0,
                  0, 0, 0, 0, 0, 0, 0, 0, 1])
    )
])
def test_matrix_elements(empty_instance, max_order, X, answer):
    dysolve = empty_instance
    dysolve._X = X
    # The basis shouldn't matter
    dysolve._basis = qeye_like(X)
    current_matrix_elements = None
    dysolve._elems = dysolve._X.transform(dysolve._basis).full().flatten()

    for _ in range(1, max_order + 1):
        current_matrix_elements = dysolve._update_matrix_elements(
            current_matrix_elements
        )
    assert np.array_equal(current_matrix_elements, answer)


@pytest.mark.parametrize("H_0", [
    sigmaz(), sigmay(), sigmaz(), qeye(2), tensor(sigmax(), sigmaz()),
    tensor(sigmax(), sigmaz()) + tensor(qeye(2), sigmay()),
    _enr_xz()
])
@pytest.mark.parametrize("t_i, t_f", [
    (0, 0.1), (0, 0.5), (0, 1), (0, 10), (0, -1),
    (-0.1, 0.1), (-0.5, 0.5), (-1, 1), (-10, 10), (1, -1)
])
def test_zeroth_order(H_0, t_i, t_f):
    # self.X and self.omega don't matter
    dysolve = DysolvePropagator(
        H_0, qeye_like(H_0), 0, options={'max_order': 0}
    )
    U = dysolve(t_f, t_i)

    exp = (-1j*H_0*(t_f - t_i)).expm()

    with CoreOptions(atol=1e-10, rtol=1e-10):
        assert U == exp


@pytest.mark.parametrize("H_0", [sigmax(), sigmay(), sigmaz()])
@pytest.mark.parametrize("X", [sigmax(), sigmay(), sigmaz()])
@pytest.mark.parametrize("t", [-0.15, -0.1, 0, 0.1, 0.15])
@pytest.mark.parametrize("omega", [0, 1, 10])
def test_2x2_propagators_single_time(H_0, X, t, omega):
    # Dysolve
    options = {'max_order': 3, 'max_dt': 0.05}
    U = dysolve_propagator(H_0, X, omega, t, options=options)

    # Qutip.solver.propagator
    def H1_coeff(t, omega):
        return np.cos(omega * t)

    H = [H_0, [X, H1_coeff]]
    args = {'omega': omega}
    prop = propagator(
        H, t, args=args, options={"atol": 1e-10, "rtol": 1e-8}
    )

    with CoreOptions(atol=1e-10, rtol=1e-6):
        assert U == prop


@pytest.mark.parametrize("H_0", [sigmay(), sigmaz()])
@pytest.mark.parametrize("X", [sigmay(), sigmaz()])
@pytest.mark.parametrize("ts", [
    [0, 0.25, 0.5],
    [0, -0.25, -0.5],
    [-0.1, 0, 0.1]
])
@pytest.mark.parametrize("omega", [0, 10])
def test_2x2_propagators_list_times(H_0, X, ts, omega):
    options = {'max_order': 3, 'max_dt': 0.01}
    Us = dysolve_propagator(H_0, X, omega, ts, options=options)

    # Qutip.solver.propagator
    def H1_coeff(t, omega):
        return np.cos(omega * t)

    H = [H_0, [X, H1_coeff]]
    args = {'omega': omega}
    props = propagator(
        H, ts, args=args, options={"atol": 1e-10, "rtol": 1e-8}
    )

    with CoreOptions(atol=1e-10, rtol=1e-6):
        assert Us == props


@pytest.mark.parametrize("H_0", [
    tensor(sigmax(), sigmaz()) + tensor(qeye(2), sigmay()),
    tensor(sigmaz(), qeye(2))
])
@pytest.mark.parametrize("X", [
    tensor(qeye(2), sigmaz()),
    tensor(sigmaz(), sigmax()) + tensor(sigmay(), qeye(2))
])
@pytest.mark.parametrize("omega", [
    5, 10
])
@pytest.mark.parametrize("t_f", [
    1, -1
])
def test_4x4_propagators_single_time(H_0, X, omega, t_f):
    options = {'max_order': 3, 'max_dt': 0.01}
    U = dysolve_propagator(H_0, X, omega, t_f, options=options)

    # Qutip.solver.propagator
    def H1_coeff(t, omega):
        return np.cos(omega * t)

    H = [H_0, [X, H1_coeff]]
    args = {'omega': omega}
    prop = propagator(
        H, t_f, args=args, options={"atol": 1e-10, "rtol": 1e-8}
    )

    with CoreOptions(atol=1e-10, rtol=1e-5):
        assert U == prop


@pytest.mark.parametrize("H_0", [
    tensor(sigmax(), sigmaz()) + tensor(qeye(2), sigmay()),
    tensor(sigmaz(), qeye(2))
])
@pytest.mark.parametrize("X", [
    tensor(qeye(2), sigmaz()),
    tensor(sigmaz(), sigmax()) + tensor(sigmay(), qeye(2))
])
@pytest.mark.parametrize("omega", [
    0, 10
])
@pytest.mark.parametrize("ts", [
    [0, 0.25, 0.5],
    [0, -0.25, -0.5],
    [-0.1, 0, 0.1]
])
def test_4x4_propagators_list_times(H_0, X, omega, ts):
    options = {'max_order': 3, 'max_dt': 0.01}
    Us = dysolve_propagator(H_0, X, omega, ts, options=options)

    # Qutip.solver.propagator
    def H1_coeff(t, omega):
        return np.cos(omega * t)

    H = [H_0, [X, H1_coeff]]
    args = {'omega': omega}
    props = propagator(
        H, ts, args=args, options={"atol": 1e-10, "rtol": 1e-8}
    )

    with CoreOptions(atol=1e-10, rtol=1e-6):
        assert Us == props


@pytest.mark.parametrize("omega", [5, 10])
@pytest.mark.parametrize("t_f", [1, -1])
def test_enr_propagators_single_time(omega, t_f):
    # reuses other test with both H_0 and X set to enr space operators
    H_0 = _enr_zz()
    X = _enr_xz()
    test_4x4_propagators_single_time(H_0, X, omega, t_f)


@pytest.mark.parametrize("omega", [0, 10])
@pytest.mark.parametrize("ts", [
    [0, 0.25, 0.5],
    [0, -0.25, -0.5],
    [-0.1, 0, 0.1]
])
def test_enr_propagators_list_times(omega, ts):
    # reuses other test with both H_0 and X set to enr space operators
    H_0 = _enr_zz()
    X = _enr_xz()
    test_4x4_propagators_list_times(H_0, X, omega, ts)


@pytest.mark.parametrize("H_0, X", [
    (
        sigmaz(), sigmax(),
    ),
    (
        tensor(sigmaz(), sigmaz()), tensor(sigmax(), sigmax()),
    ),
    (
        tensor(sigmaz(), sigmaz(), sigmaz()),
        tensor(sigmax(), sigmax(), sigmax())
    ),
    (
        tensor(sigmaz(), sigmaz(), sigmaz(), sigmaz()),
        tensor(sigmax(), sigmax(), sigmax(), sigmax())
    ),
    (
        _enr_zz(), _enr_xx()
    )
])
def test_dims(H_0, X):
    dysolve = DysolvePropagator(
        H_0, X, 1, {'max_order': 0}
    )
    U = dysolve(0.001)
    assert (dysolve._H_0.dims == dysolve._X.dims == U.dims)
