import pytest
import pickle
import qutip
import numpy as np
from qutip.solver.sesolve import sesolve, SESolver
from qutip.solver.krylovsolve import krylovsolve
from qutip.solver.solver_base import Solver

# Deactivate warning for test without cython
from qutip.core.coefficient import WARN_MISSING_MODULE
WARN_MISSING_MODULE[0] = 0


all_ode_method = [
    method for method, integrator in SESolver.avail_integrators().items()
    if integrator.support_time_dependant
]

def _analytic(t, alpha):
    return ((1 - np.exp(-alpha * t)) / alpha)


class TestSeSolve():
    H0 = 0.2 * np.pi * qutip.sigmaz()
    H1 = np.pi * qutip.sigmax()
    tlist = np.linspace(0, 20, 200)
    args = {'alpha': 0.5}
    w_a = 0.35
    a = 0.5

    @pytest.mark.parametrize(['unitary_op'], [
        pytest.param(None, id="state"),
        pytest.param(qutip.qeye(2), id="unitary"),
    ])
    @pytest.mark.parametrize(['H', 'analytical'], [
        pytest.param(H1, lambda t, _: t, id='const_H'),
        pytest.param(lambda t, alpha: (np.pi * qutip.sigmax()
                                       * np.exp(-alpha * t)),
                     _analytic, id='func_H'),
        pytest.param([[H1, lambda t, args: np.exp(-args['alpha'] * t)]],
                     _analytic, id='list_func_H'),
        pytest.param([[H1, 'exp(-alpha*t)']],
                     _analytic, id='list_str_H'),
        pytest.param([[H1, np.exp(-args['alpha'] * tlist)]],
                     _analytic, id='list_array_H'),
        pytest.param(qutip.QobjEvo([[H1, 'exp(-alpha*t)']], args=args),
                     _analytic, id='QobjEvo_H'),
    ])
    def test_sesolve(self, H, analytical, unitary_op):
        """
        Compare integrated evolution with analytical result
        If U0 is not None then operator evo is checked
        Otherwise state evo
        """
        tol = 5e-3
        psi0 = qutip.basis(2, 0)
        options = {"progress_bar": None}

        if unitary_op is None:
            output = sesolve(H, psi0, self.tlist,
                             [qutip.sigmax(), qutip.sigmay(), qutip.sigmaz()],
                             args=self.args, options=options)
            sx, sy, sz = output.expect[0], output.expect[1], output.expect[2]
        else:
            output = sesolve(H, unitary_op, self.tlist, args=self.args,
                             options=options)
            sx = [qutip.expect(qutip.sigmax(), U * psi0)
                  for U in output.states]
            sy = [qutip.expect(qutip.sigmay(), U * psi0)
                  for U in output.states]
            sz = [qutip.expect(qutip.sigmaz(), U * psi0)
                  for U in output.states]

        sx_analytic = np.zeros(np.shape(self.tlist))
        sy_analytic = np.array(
            [-np.sin(2 * np.pi * analytical(t, self.args['alpha']))
             for t in self.tlist]
        )
        sz_analytic = np.array(
            [np.cos(2 * np.pi * analytical(t, self.args['alpha']))
             for t in self.tlist]
        )

        np.testing.assert_allclose(sx, sx_analytic, atol=tol)
        np.testing.assert_allclose(sy, sy_analytic, atol=tol)
        np.testing.assert_allclose(sz, sz_analytic, atol=tol)

    @pytest.mark.parametrize(['state_type'], [
        pytest.param("ket", id="ket"),
        pytest.param("unitary", id="unitary"),
    ])
    def test_sesolve_normalization(self, state_type):
        # non-hermitean H causes state to evolve non-unitarily
        H = qutip.Qobj([[1, -0.1j], [-0.1j, 1]])
        psi0 = qutip.basis(2, 0)
        options = {"normalize_output": True, "progress_bar": None}

        if state_type == "ket":
            output = sesolve(H, psi0, self.tlist, e_ops=[], options=options)
            norms = [state.norm() for state in output.states]
            np.testing.assert_allclose(
                norms, [1.0 for _ in self.tlist], atol=1e-15,
            )
        else:
            # evolution of unitaries should not be normalized
            U = qutip.qeye(2)
            output = sesolve(H, U, self.tlist, e_ops=[], options=options)
            norms = [state.norm() for state in output.states]
            assert all(norm > 2 for norm in norms[1:])

    @pytest.mark.parametrize(['unitary_op'], [
        pytest.param(None, id="state"),
        pytest.param(qutip.qeye(2), id="unitary"),
    ])
    @pytest.mark.parametrize('method', all_ode_method, ids=all_ode_method)
    def test_sesolve_method(self, method, unitary_op):
        """
        Compare integrated evolution with analytical result
        If U0 is not None then operator evo is checked
        Otherwise state evo
        """
        tol = 5e-3
        psi0 = qutip.basis(2, 0)
        options = {"method": method, "progress_bar": None}
        H = [[self.H1, 'exp(-alpha*t)']]

        if unitary_op is None:
            output = sesolve(H, psi0, self.tlist,
                             [qutip.sigmax(), qutip.sigmay(), qutip.sigmaz()],
                             args=self.args, options=options)
            sx, sy, sz = output.expect[0], output.expect[1], output.expect[2]
        else:
            output = sesolve(H, unitary_op, self.tlist,
                             args=self.args, options=options)
            sx = [qutip.expect(qutip.sigmax(), U * psi0)
                  for U in output.states]
            sy = [qutip.expect(qutip.sigmay(), U * psi0)
                  for U in output.states]
            sz = [qutip.expect(qutip.sigmaz(), U * psi0)
                  for U in output.states]

        sx_analytic = np.zeros(np.shape(self.tlist))
        sy_analytic = np.array(
            [np.sin(-2*np.pi * _analytic(t, self.args['alpha']))
             for t in self.tlist]
        )
        sz_analytic = np.array(
            [np.cos(2*np.pi *_analytic(t, self.args['alpha']))
             for t in self.tlist]
        )

        np.testing.assert_allclose(sx, sx_analytic, atol=tol)
        np.testing.assert_allclose(sy, sy_analytic, atol=tol)
        np.testing.assert_allclose(sz, sz_analytic, atol=tol)

    @pytest.mark.parametrize('normalize', [True, False],
                             ids=['Normalized', ''])
    @pytest.mark.parametrize(['H', 'args'],
        [pytest.param(H0 + H1, {}, id='const_H'),
         pytest.param(lambda t, a, w_a: (
             a * t * 0.2 * np.pi * qutip.sigmaz() +
             np.cos(w_a * t) * np.pi * qutip.sigmax()
         ), {'a':a, 'w_a':w_a}, id='func_H'),
         pytest.param([
             [H0, lambda t, args: args['a']*t],
             [H1, lambda t, args: np.cos(args['w_a']*t)]
         ], {'a':a, 'w_a':w_a}, id='list_func_H'),
         pytest.param([H0, [H1, 'cos(w_a*t)']], {'w_a':w_a}, id='list_str_H'),
    ])
    def test_compare_evolution(self, H, normalize, args, tol=5e-5):
        """
        Compare integrated evolution of unitary operator with state evo
        """
        psi0 = qutip.basis(2, 0)
        U0 = qutip.qeye(2)

        options = {
            "store_states": True,
            "normalize_output": normalize,
            "progress_bar": None
        }
        out_s = sesolve(H, psi0, self.tlist, [qutip.sigmax(),
                                              qutip.sigmay(),
                                              qutip.sigmaz()],
                        options=options, args=args)
        xs, ys, zs = out_s.expect[0], out_s.expect[1], out_s.expect[2]
        xss = [qutip.expect(qutip.sigmax(), U) for U in out_s.states]
        yss = [qutip.expect(qutip.sigmay(), U) for U in out_s.states]
        zss = [qutip.expect(qutip.sigmaz(), U) for U in out_s.states]

        np.testing.assert_allclose(xs, xss, atol=tol)
        np.testing.assert_allclose(ys, yss, atol=tol)
        np.testing.assert_allclose(zs, zss, atol=tol)

        if normalize:
            # propagator evolution is not normalized (yet?)
            tol = 5e-4
        out_u = sesolve(H, U0, self.tlist, options=options, args=args)
        xu = [qutip.expect(qutip.sigmax(), U * psi0) for U in out_u.states]
        yu = [qutip.expect(qutip.sigmay(), U * psi0) for U in out_u.states]
        zu = [qutip.expect(qutip.sigmaz(), U * psi0) for U in out_u.states]

        np.testing.assert_allclose(xs, xu, atol=tol)
        np.testing.assert_allclose(ys, yu, atol=tol)
        np.testing.assert_allclose(zs, zu, atol=tol)

    def test_sesolver_args(self):
        options = {"progress_bar": None}
        solver_obj = SESolver(qutip.QobjEvo([self.H0, [self.H1,'a']],
                                            args={'a': 1}),
                              options=options)
        res = solver_obj.run(qutip.basis(2,1), [0, 1, 2, 3],
                             e_ops=[qutip.num(2)], args={'a':0})
        np.testing.assert_allclose(res.expect[0], 1)

    def test_sesolver_pickling(self):
        e_ops = [qutip.sigmax(), qutip.sigmay(), qutip.sigmaz()]
        options = {"progress_bar": None}
        solver_obj = SESolver(self.H0 + self.H1,
                              options=options)
        solver_copy = pickle.loads(pickle.dumps(solver_obj))
        sx, sy, sz = solver_obj.run(qutip.basis(2,1), [0, 1, 2, 3],
                                    e_ops=e_ops).expect
        csx, csy, csz = solver_copy.run(qutip.basis(2,1), [0, 1, 2, 3],
                                       e_ops=e_ops).expect
        np.testing.assert_allclose(sx, csx)
        np.testing.assert_allclose(sy, csy)
        np.testing.assert_allclose(sz, csz)

    @pytest.mark.parametrize('method', all_ode_method, ids=all_ode_method)
    def test_sesolver_stepping(self, method):
        options = {
            "method": method,
            "atol": 1e-7,
            "rtol": 1e-8,
            "progress_bar": None
        }
        solver_obj = SESolver(
            qutip.QobjEvo([self.H1, lambda t, a: a], args={"a":0.25}),
            options=options
        )
        solver_obj.start(qutip.basis(2,0), 0)
        sr2 = -(2**.5)/2
        state = solver_obj.step(1)
        np.testing.assert_allclose(qutip.expect(qutip.sigmax(), state), 0.,
                                   atol=2e-6)
        np.testing.assert_allclose(qutip.expect(qutip.sigmay(), state), -1,
                                   atol=2e-6)
        np.testing.assert_allclose(qutip.expect(qutip.sigmaz(), state), 0.,
                                   atol=2e-6)

        state = solver_obj.step(2, args={"a":0.125})
        np.testing.assert_allclose(qutip.expect(qutip.sigmax(), state), 0.,
                                   atol=2e-6)
        np.testing.assert_allclose(qutip.expect(qutip.sigmay(), state), sr2,
                                   atol=2e-6)
        np.testing.assert_allclose(qutip.expect(qutip.sigmaz(), state), sr2,
                                   atol=2e-6)

        solver_obj.options = {
            "method": "adams",
            "atol": 1e-7,
            "rtol": 1e-8,
            "progress_bar": None
        }
        state = solver_obj.step(3, args={"a":0})
        np.testing.assert_allclose(qutip.expect(qutip.sigmax(), state), 0.,
                                   atol=2e-6)
        np.testing.assert_allclose(qutip.expect(qutip.sigmay(), state), sr2,
                                   atol=2e-6)
        np.testing.assert_allclose(qutip.expect(qutip.sigmaz(), state), sr2,
                                   atol=2e-6)


def test_sesolve_bad_H():
    with pytest.raises(TypeError):
        SESolver(np.eye(3))
    with pytest.raises(ValueError):
        SESolver(qutip.basis(3,1))


def test_sesolve_bad_state():
    solver = SESolver(qutip.qeye(4))
    with pytest.raises(TypeError):
        solver.start(qutip.basis(4,1).dag(), 0)
    with pytest.raises(TypeError):
        solver.start(qutip.basis(2,1) & qutip.basis(2,0), 0)


def test_sesolve_step_no_start():
    solver = SESolver(qutip.qeye(4))
    with pytest.raises(RuntimeError):
        solver.step(1)


@pytest.mark.parametrize("always_compute_step", [True, False])
def test_krylovsolve(always_compute_step):
    H = qutip.tensor([qutip.rand_herm(2) for _ in range(8)])
    psi0 = qutip.basis([2]*8, [1]*8)
    e_op = qutip.num(256)
    e_op.dims = H.dims
    tlist = np.linspace(0, 1, 11)
    ref = sesolve(H, psi0, tlist, e_ops=[e_op]).expect[0]
    options = {"always_compute_step": always_compute_step}
    krylov_sol = krylovsolve(H, psi0, tlist, 20, e_ops=[e_op], options=options)
    np.testing.assert_allclose(ref, krylov_sol.expect[0])


def test_krylovsolve_error():
    H = qutip.rand_herm(256, density=0.2)
    psi0 = qutip.basis([256], [255])
    tlist = np.linspace(0, 1, 11)
    options = {"min_step": 1e10}
    with pytest.raises(ValueError) as err:
        krylovsolve(H, psi0, tlist, 20, options=options)
    assert "error with the minimum step" in str(err.value)


def test_feedback():

    def f(t, A, qobj=None):
        return (A-2.)

    N = 4
    tol = 1e-14
    psi0 = qutip.basis(N, N-1)
    a = qutip.destroy(N)
    H = qutip.QobjEvo(
        [qutip.num(N), [a+a.dag(), f]],
        args={"A": SESolver.ExpectFeedback(qutip.num(N), default=3.)}
    )
    solver = qutip.SESolver(H)
    result = solver.run(psi0, np.linspace(0, 30, 301), e_ops=[qutip.num(N)])
    assert np.all(result.expect[0] > 2 - tol)