import numpy as np
from qutip import (
    sigmax, sigmay, sigmaz,  sigmap, sigmam,
    rand_ket, num, destroy,
    mesolve, expect, sesolve,
    Qobj, QobjEvo, coefficient
 )

from qutip.solver.floquet import (
    FloquetBasis, floquet_tensor, fmmesolve, FMESolver,
    _floquet_delta_tensor, _floquet_X_matrices, fsesolve
)
import pytest


def _convert_c_ops(c_op_fmmesolve, noise_spectrum, vp, ep):
    """
    Convert e_ops for fmmesolve to mesolve
    """
    c_op_mesolve = []
    N = len(vp)
    for i in range(N):
        for j in range(N):
            if i != j:
                # caclculate the rate
                gamma = 2 * np.pi * c_op_fmmesolve.matrix_element(
                    vp[j], vp[i]) * c_op_fmmesolve.matrix_element(
                    vp[i], vp[j]) * noise_spectrum(ep[j] - ep[i])

                # add c_op for mesolve
                c_op_mesolve.append(
                    np.sqrt(gamma) * (vp[i] * vp[j].dag())
                )
    return c_op_mesolve


class TestFloquet:
    """
    A test class for the QuTiP functions for Floquet formalism.
    """

    def testFloquetBasis(self):
        N = 10
        a = destroy(N)
        H = num(N) + (a+a.dag()) * coefficient(lambda t: np.cos(t))
        T = 2 * np.pi
        floquet_basis = FloquetBasis(H, T)
        psi0 = rand_ket(N)
        tlist = np.linspace(0, 10, 11)
        floquet_psi0 = floquet_basis.to_floquet_basis(psi0)
        states = sesolve(H, psi0, tlist).states
        for t, state in zip(tlist, states):
            from_floquet = floquet_basis.from_floquet_basis(floquet_psi0, t)
            assert state.overlap(from_floquet) == pytest.approx(1., abs=8e-5)

    def testFloquetUnitary(self):
        N = 10
        a = destroy(N)
        H = num(N) + (a+a.dag()) * coefficient(lambda t: np.cos(t))
        T = 2 * np.pi
        psi0 = rand_ket(N)
        tlist = np.linspace(0, 10, 11)
        states_se = sesolve(H, psi0, tlist).states
        states_fse = fsesolve(H, psi0, tlist, T=T).states
        for state_se, state_fse in zip(states_se, states_fse):
            assert state_se.overlap(state_fse) == pytest.approx(1., abs=5e-5)


    def testFloquetMasterEquation1(self):
        """
        Test Floquet-Markov Master Equation for a driven two-level system
        without dissipation.
        """
        delta = 1.0 * 2 * np.pi
        eps0 = 1.0 * 2 * np.pi
        A = 0.5 * 2 * np.pi
        omega = np.sqrt(delta ** 2 + eps0 ** 2)
        T = (2 * np.pi) / omega
        tlist = np.linspace(0.0, 2 * T, 101)
        psi0 = rand_ket(2)
        H0 = - eps0 / 2.0 * sigmaz() - delta / 2.0 * sigmax()
        H1 = A / 2.0 * sigmax()
        args = {'w': omega}
        H = [H0, [H1, lambda t, args: np.sin(args['w'] * t)]]
        e_ops = [num(2)]
        gamma1 = 0

        # Collapse operator for Floquet-Markov Master Equation
        c_op_fmmesolve = sigmax()

        # Collapse operators for Lindblad Master Equation
        def spectrum(omega):
            return (omega > 0) * omega * 0.5 * gamma1 / (2 * np.pi)

        ep, vp = H0.eigenstates()
        op0 = vp[0]*vp[0].dag()
        op1 = vp[1]*vp[1].dag()

        c_op_mesolve = _convert_c_ops(c_op_fmmesolve, spectrum, vp, ep)

        # Solve the floquet-markov master equation
        p_ex = fmmesolve(H, psi0, tlist, [c_op_fmmesolve], [num(2)],
                         [spectrum], T, args=args).expect[0]

        # Compare with mesolve
        p_ex_ref = mesolve(
            H, psi0, tlist, c_op_mesolve, [num(2)], args
        ).expect[0]

        np.testing.assert_allclose(np.real(p_ex), np.real(p_ex_ref), atol=1e-4)

    def testFloquetMasterEquation2(self):
        """
        Test Floquet-Markov Master Equation for a two-level system
        subject to dissipation.
        """
        delta = 1.0 * 2 * np.pi
        eps0 = 1.0 * 2 * np.pi
        A = 0.0 * 2 * np.pi
        omega = np.sqrt(delta ** 2 + eps0 ** 2)
        T = (2 * np.pi) / omega
        tlist = np.linspace(0.0, 2 * T, 101)
        psi0 = rand_ket(2)
        H0 = - eps0 / 2.0 * sigmaz() - delta / 2.0 * sigmax()
        H1 = A / 2.0 * sigmax()
        args = {'w': omega}
        H = QobjEvo([H0, [H1, lambda t, w: np.sin(w * t)]],
                    args=args)
        e_ops = [num(2)]
        gamma1 = 1

        # Collapse operator for Floquet-Markov Master Equation
        c_op_fmmesolve = sigmax()

        # Collapse operator for Lindblad Master Equation
        def spectrum(omega):
            return (omega > 0) * omega * 0.5 * gamma1 / (2 * np.pi)

        ep, vp = H0.eigenstates()
        op0 = vp[0]*vp[0].dag()
        op1 = vp[1]*vp[1].dag()

        c_op_mesolve = _convert_c_ops(c_op_fmmesolve, spectrum, vp, ep)

        # Solve the floquet-markov master equation
        p_ex = fmmesolve(
            H, psi0, tlist, [c_op_fmmesolve], [num(2)],
            [spectrum], T, args=args
        ).expect[0]

        # Compare with mesolve
        p_ex_ref = mesolve(
            H, psi0, tlist, c_op_mesolve, [num(2)], args
        ).expect[0]

        np.testing.assert_allclose(np.real(p_ex), np.real(p_ex_ref), atol=1e-4)

    @pytest.mark.parametrize("kmax", [5, 25, 100])
    def testFloquetMasterEquation3(self, kmax):
        """
        Test Floquet-Markov Master Equation for a two-level system
        subject to dissipation with internal transform of fmmesolve
        """
        delta = 1.0 * 2 * np.pi
        eps0 = 1.0 * 2 * np.pi
        A = 0.0 * 2 * np.pi
        omega = np.sqrt(delta ** 2 + eps0 ** 2)
        T = (2 * np.pi) / omega
        tlist = np.linspace(0.0, 2 * T, 101)
        psi0 = rand_ket(2)
        H0 = - eps0 / 2.0 * sigmaz() - delta / 2.0 * sigmax()
        H1 = A / 2.0 * sigmax()
        args = {'w': omega}
        H = QobjEvo([H0, [H1, lambda t, w: np.sin(w * t)]],
                    args=args)
        e_ops = [num(2)]
        gamma1 = 1

        # Collapse operator for Floquet-Markov Master Equation
        c_op_fmmesolve = sigmax()

        # Collapse operator for Lindblad Master Equation
        def spectrum(omega):
            return (omega > 0) * 0.5 * gamma1 * omega / (2 * np.pi)

        ep, vp = H0.eigenstates()
        op0 = vp[0]*vp[0].dag()
        op1 = vp[1]*vp[1].dag()

        c_op_mesolve = _convert_c_ops(c_op_fmmesolve, spectrum, vp, ep)

        # Solve the floquet-markov master equation
        floquet_basis = FloquetBasis(H, T)
        solver = FMESolver(
            floquet_basis, [(c_op_fmmesolve, spectrum)], kmax=kmax
        )
        solver.start(psi0, tlist[0])
        p_ex = [expect(num(2), solver.step(t)) for t in tlist]

        # Compare with mesolve
        output2 = mesolve(H, psi0, tlist, c_op_mesolve, [num(2)], args)
        p_ex_ref = output2.expect[0]

        np.testing.assert_allclose(np.real(p_ex), np.real(p_ex_ref),
                                   atol=5 * 1e-4)

    def testFloquetMasterEquation_multiple_coupling(self):
        """
        Test Floquet-Markov Master Equation for a two-level system
        subject to dissipation with multiple coupling operators
        """
        delta = 1.0 * 2 * np.pi
        eps0 = 1.0 * 2 * np.pi
        A = 0.0 * 2 * np.pi
        omega = np.sqrt(delta ** 2 + eps0 ** 2)
        T = (2 * np.pi) / omega
        tlist = np.linspace(0.0, 2 * T, 101)
        psi0 = rand_ket(2)
        H0 = - eps0 / 2.0 * sigmaz() - delta / 2.0 * sigmax()
        H1 = A / 2.0 * sigmax()
        args = {'w': omega}
        H = QobjEvo([H0, [H1, lambda t, w: np.sin(w * t)]],
                    args=args)
        e_ops = [num(2)]
        gamma1 = 1

        # Collapse operator for Floquet-Markov Master Equation
        c_ops_fmmesolve = [sigmax(), sigmay()]

        # Collapse operator for Lindblad Master Equation
        def noise_spectrum1(omega):
            return (omega > 0) * 0.5 * gamma1 * omega/(2*np.pi)

        def noise_spectrum2(omega):
            return (omega > 0) * 0.5 * gamma1 / (2 * np.pi)

        noise_spectra = [noise_spectrum1, noise_spectrum2]

        ep, vp = H0.eigenstates()
        op0 = vp[0]*vp[0].dag()
        op1 = vp[1]*vp[1].dag()

        c_op_mesolve = []

        # Convert the c_ops for fmmesolve to c_ops for mesolve
        for c_op_fmmesolve, noise_spectrum in zip(c_ops_fmmesolve,
                                                   noise_spectra):
            c_op_mesolve += _convert_c_ops(
                c_op_fmmesolve, noise_spectrum, vp, ep
            )

        # Solve the floquet-markov master equation
        output1 = fmmesolve(
            H, psi0, tlist, c_ops_fmmesolve, e_ops, noise_spectra, T
        )
        p_ex = output1.expect[0]
        # Compare with mesolve
        output2 = mesolve(H, psi0, tlist, c_op_mesolve, e_ops, args)
        p_ex_ref = output2.expect[0]

        np.testing.assert_allclose(np.real(p_ex), np.real(p_ex_ref), atol=1e-4)

    def testFloquetRates(self):
        """
        Compare rate transition and frequency transitions to analytical
        results for a driven two-level system, for different drive amplitudes.
        """
        # Parameters
        wq = 5.4 * 2 * np.pi
        wd = 6.7 * 2 * np.pi
        delta = wq - wd
        T = 2 * np.pi / wd
        tlist = np.linspace(0.0, 2 * T, 101)
        array_A = np.linspace(0.001, 3 * 2 * np.pi, 10, endpoint=False)
        H0 = Qobj(wq / 2 * sigmaz())
        args = {'wd': wd}
        c_ops = sigmax()
        gamma1 = 1
        kmax = 5

        array_ana_E0 = [-np.sqrt((delta / 2)**2 + a**2) for a in array_A]
        array_ana_E1 = [np.sqrt((delta / 2)**2 + a**2) for a in array_A]
        array_ana_delta = [
            2 * np.sqrt((delta / 2)**2 + a**2)
            for a in array_A
        ]

        def noise_spectrum(omega):
            return (omega > 0) * 0.5 * gamma1 * omega/(2*np.pi)

        idx = 0
        for a in array_A:
            # Hamiltonian
            H1_p = a * sigmap()
            H1_m = a * sigmam()
            H = QobjEvo(
                [H0, [H1_p, lambda t, wd: np.exp(-1j * wd * t)],
                [H1_m, lambda t, wd: np.exp(1j * wd * t)]], args=args
            )

            floquet_basis = FloquetBasis(H, T)
            DeltaMatrix = _floquet_delta_tensor(floquet_basis.e_quasi, kmax, T)
            X = _floquet_X_matrices(floquet_basis, [c_ops], kmax)

            # Check energies
            deltas = np.ndarray.flatten(DeltaMatrix)

            # deltas and array_ana_delta have at least 1 value in common.
            assert (min(abs(deltas - array_ana_delta[idx])) < 1e-4)

            # Check matrix elements
            Xs = np.concatenate(
                [X[0][k].to_array() for k in range(-kmax, kmax+1)]
            ).flatten()

            normPlus = np.sqrt(a**2 + (array_ana_E1[idx] - delta / 2)**2)
            normMinus = np.sqrt(a**2 + (array_ana_E0[idx] - delta / 2)**2)

            Xpp_p1 = (a / normPlus**2) * (array_ana_E1[idx] - delta / 2)
            assert (min(abs(Xs - Xpp_p1)) < 1e-4)
            Xpp_m1 = (a / normPlus**2) * (array_ana_E1[idx] - delta / 2)
            assert (min(abs(Xs - Xpp_m1)) < 1e-4)
            Xmm_p1 = (a / normMinus**2) * (array_ana_E0[idx] - delta / 2)
            assert (min(abs(Xs - Xmm_p1)) < 1e-4)
            Xmm_m1 = (a / normMinus**2) * (array_ana_E0[idx] - delta / 2)
            assert (min(abs(Xs - Xmm_m1)) < 1e-4)
            Xpm_p1 = (a / (normMinus * normPlus)
                      * (array_ana_E0[idx]- delta / 2))
            assert (min(abs(Xs - Xpm_p1)) < 1e-4)
            Xpm_m1 = (a / (normMinus * normPlus)
                      * (array_ana_E1[idx] - delta / 2))
            assert (min(abs(Xs - Xpm_m1)) < 1e-4)
            idx += 1


def test_fsesolve_fallback():
    H = [sigmaz(), lambda t: np.sin(t * 2 * np.pi)]
    psi0 = rand_ket(2)
    ffstate = fmmesolve(H, psi0, [0, 1], T=1.).final_state
    fstate = sesolve(H, psi0, [0, 1]).final_state
    assert (ffstate - fstate).norm() < 1e-5