import pytest
import functools
from itertools import product
import numpy as np
from scipy.integrate import trapezoid
import qutip

pytestmark = [pytest.mark.usefixtures("in_temporary_directory")]

_equivalence_dimension = 15
_equivalence_fock = qutip.fock(_equivalence_dimension, 1)
_equivalence_coherent = qutip.coherent_dm(_equivalence_dimension, 2)


@pytest.mark.filterwarnings("ignore::FutureWarning")
@pytest.mark.parametrize(["solver", "start"], [
    pytest.param("es", _equivalence_coherent, id="es"),
    pytest.param("es", None, id="es-steady state"),
])
def test_correlation_solver_equivalence(solver, start):
    """
    Test that all of the correlation solvers give the same results for a given
    system.
    """
    a = qutip.destroy(_equivalence_dimension)
    x = ( a + a.dag() ) / np.sqrt(2)
    H = a.dag() * a
    G1 = 0.75
    n_th = 2
    c_ops = [np.sqrt(G1 * (n_th+1)) * a,
             np.sqrt(G1 * n_th) * a.dag()]
    times = np.linspace(0, 3, 61)
    # Massively relax the tolerance for the Monte-Carlo approach to avoid a
    # long simulation time.
    tol = 0.25 if solver == "mc" else 1e-4
    # We use the master equation version as a base, but it doesn't actually
    # matter - if all the tests fail, it implies that the "me" solver might be
    # broken, whereas if only one fails, then it implies that only that one is
    # broken.  We test that all solvers are equivalent by transitive equality
    # to the "me" solver.
    base = qutip.correlation_2op_2t(H, start, None, times, c_ops,
                                    a.dag(), a, solver="me")
    cmp = qutip.correlation_2op_2t(H, start, None, times, c_ops,
                                   a.dag(), a, solver=solver)
    np.testing.assert_allclose(base, cmp, atol=tol)


def _spectrum_wrapper(solver):
    frequencies = 2*np.pi * np.linspace(0.5, 1.5, 101)
    @functools.wraps(qutip.spectrum)
    def out(H, c_ops, a, b):
        return (qutip.spectrum(H, frequencies, c_ops, a, b, solver=solver),
                frequencies)
    return out


def _spectrum_fft(H, c_ops, a, b):
    times = np.linspace(0, 100, 2500)
    correlation = qutip.correlation_2op_1t(H, None, times, c_ops, a, b)
    frequencies, spectrum = qutip.spectrum_correlation_fft(times, correlation)
    return spectrum, frequencies


@pytest.mark.parametrize("spectrum", [
    pytest.param(_spectrum_fft, id="fft"),
    pytest.param(_spectrum_wrapper("es"), id="es"),
    pytest.param(_spectrum_wrapper("pi"), id="pi"),
    pytest.param(_spectrum_wrapper("solve"), id="solve"),
])
def test_spectrum_solver_equivalence_to_es(spectrum):
    """Test equivalence of the spectrum solvers to the base "es" method."""
    # Jaynes--Cummings model.
    dimension = 4
    wc = wa = 1.0 * 2*np.pi
    g = 0.1 * 2*np.pi
    kappa = 0.75
    gamma = 0.25
    n_th = 0.01

    a = qutip.tensor(qutip.destroy(dimension), qutip.qeye(2))
    sm = qutip.tensor(qutip.qeye(dimension), qutip.sigmam())
    H = wc*a.dag()*a + wa*sm*sm.dag() + g*(a.dag()*sm.dag() + a*sm)
    c_ops = [np.sqrt(kappa * (n_th+1)) * a,
             np.sqrt(kappa * n_th) * a.dag(),
             np.sqrt(gamma) * sm.dag()]

    test, frequencies = spectrum(H, c_ops, a.dag(), a)
    base = qutip.spectrum(H, frequencies, c_ops, a.dag(), a, solver="es")
    np.testing.assert_allclose(base, test, atol=1e-3)


def _trapz_2d(z, xy):
    """2D trapezium-method integration assuming a square grid."""
    dx = xy[1] - xy[0]
    return dx*dx * trapezoid(trapezoid(z, axis=0))


def _n_correlation(times, n):
    """
    Numerical integration of the correlation function given an array of
    expectation values.
    """
    return np.array([[n[t] * n[t+tau] for tau in range(times.shape[0])]
                     for t in range(times.shape[0])])


def _coefficient_function(t, args):
    t_off, tp = args['t_off'], args['tp']
    return np.exp(-(t-t_off)*(t-t_off) / (2*tp*tp))


_coefficient_string = "exp(-(t-t_off)**2 / (2 * tp*tp))"


def _h_qobj_function(t, args):
    return args['H0'] * _coefficient_function(t, args)


# 2LS and 3LS stand for two- and three-level system respectively.

_2ls_args = {'H0': 2*qutip.sigmax(), 't_off': 1, 'tp': 0.5}
_2ls_times = np.linspace(0, 5, 51)
_3ls_args = {'t_off': 2, 'tp': 1}
_3ls_times = np.linspace(0, 6, 20)


def _2ls_g2_0(H, c_ops):
    sp = qutip.sigmap()
    start = qutip.basis(2, 0)
    times = _2ls_times
    H = qutip.QobjEvo(H, args=_2ls_args, tlist=times)
    correlation = qutip.correlation_3op_2t(H, start, times, times, [sp],
                                           sp.dag(), sp.dag()*sp, sp,
                                           args=_2ls_args)
    n_expectation = qutip.mesolve(H, start, times, [sp] + c_ops,
                                  e_ops=[qutip.num(2)],
                                  args=_2ls_args).expect[0]
    integral_correlation = _trapz_2d(np.real(correlation), times)
    integral_n_expectation = trapezoid(n_expectation, times)
    # Factor of two from negative time correlations.
    return 2 * integral_correlation / integral_n_expectation**2


@pytest.fixture(params=[
    pytest.param(_coefficient_string, id="string"),
    pytest.param(_coefficient_function(_2ls_times, _2ls_args), id="numpy"),
    pytest.param(_coefficient_function, id="function"),
])
def dependence_2ls(request):
    return request.param


class TestTimeDependence:
    """
    Test correlations with time-dependent operators using a two-level system
    (2LS) or a three-level system (3LS).
    """
    def test_varying_coefficient_hamiltonian_2ls(self, dependence_2ls):
        H = [[_2ls_args['H0'], dependence_2ls]]
        assert abs(_2ls_g2_0(H, []) - 0.575) < 1e-2

    def test_hamiltonian_from_function_2ls(self):
        H = _h_qobj_function
        assert abs(_2ls_g2_0(H, []) - 0.575) < 1e-2

    @pytest.mark.slow
    def test_varying_coefficient_hamiltonian_c_ops_2ls(self, dependence_2ls):
        H = [[_2ls_args['H0'], dependence_2ls]]
        c_ops = [[2*qutip.sigmam()*qutip.sigmap(), dependence_2ls]]
        assert abs(_2ls_g2_0(H, c_ops) - 0.824) < 1e-2

    @pytest.mark.slow
    @pytest.mark.parametrize("dependence_3ls", [
        pytest.param(_coefficient_string, id="string"),
        pytest.param(_coefficient_function(_3ls_times, _3ls_args), id="numpy"),
        pytest.param(_coefficient_function, id="function"),
    ])
    def test_coefficient_c_ops_3ls(self, dependence_3ls):
        # Calculate zero-delay HOM cross-correlation for incoherently pumped
        # three-level system, g2_ab[0] with gamma = 1.
        dimension = 3
        H = qutip.qzero(dimension)
        start = qutip.basis(dimension, 2)
        times = _3ls_times
        project_0_1 = qutip.projection(dimension, 0, 1)
        project_1_2 = qutip.projection(dimension, 1, 2)
        population_1 = qutip.projection(dimension, 1, 1)
        # Define the pi pulse to be when 99% of the population is transferred.
        rabi = np.sqrt(-np.log(0.01) / (_3ls_args['tp']*np.sqrt(np.pi)))
        coeff_3ls = qutip.coefficient(dependence_3ls, tlist=times,
                                      args=_3ls_args)
        c_ops = [project_0_1, [rabi*project_1_2, coeff_3ls]]
        forwards = qutip.correlation_2op_2t(H, start, times, times, c_ops,
                                            project_0_1.dag(), project_0_1,
                                            args=_3ls_args)
        backwards = qutip.correlation_2op_2t(H, start, times, times, c_ops,
                                             project_0_1.dag(), project_0_1,
                                             args=_3ls_args, reverse=True)
        times2 = np.concatenate([times, times[1:] + times[-1]])
        n_expect = qutip.mesolve(H, start, times2, c_ops, args=_3ls_args,
                                 e_ops=[population_1]).expect[0]
        correlation_ab = -forwards*backwards + _n_correlation(times, n_expect)
        g2_ab_0 = _trapz_2d(np.real(correlation_ab), times)
        assert abs(g2_ab_0 - 0.185) < 1e-2


def _step(t):
    return np.arctan(t)/np.pi + 0.5


def test_hamiltonian_order_unimportant():
    # Testing for regression on issue 1048.
    sp = qutip.sigmap()
    H = [[qutip.sigmax(), lambda t: _step(t-2)],
         [qutip.qeye(2), lambda t: _step(-(t-2))]]
    start = qutip.basis(2, 0)
    times = np.linspace(0, 5, 6)
    forwards = qutip.correlation_2op_2t(H, start, times, times, [sp],
                                        sp.dag(), sp)
    backwards = qutip.correlation_2op_2t(H[::-1], start, times, times, [sp],
                                         sp.dag(), sp)
    np.testing.assert_allclose(forwards, backwards, atol=1e-6)


@pytest.mark.parametrize(['solver', 'state'], [
    pytest.param('me', _equivalence_fock, id="me-ket"),
    pytest.param('me', _equivalence_coherent, id="me-dm"),
    pytest.param('me', None, id="me-steady"),
    pytest.param('es', _equivalence_fock, id="es-ket"),
    pytest.param('es', _equivalence_coherent, id="es-dm"),
    pytest.param('es', None, id="es-steady"),
])
@pytest.mark.parametrize("is_e_op_hermitian", [True, False],
                         ids=["hermitian", "nonhermitian"])
@pytest.mark.parametrize("w", [1, 2])
@pytest.mark.parametrize("gamma", [1, 10])
def test_correlation_2op_1t_known_cases(solver,
                                        state,
                                        is_e_op_hermitian,
                                        w,
                                        gamma,
                                       ):
    """This test compares the output correlation_2op_1 solution to an analytical
    solution."""

    a = qutip.destroy(_equivalence_dimension)
    x = (a + a.dag())/np.sqrt(2)

    H = w * a.dag() * a

    a_op = x if is_e_op_hermitian else a
    b_op = x if is_e_op_hermitian else a.dag()
    c_ops = [np.sqrt(gamma) * a]

    times = np.linspace(0, 1, 30)

    # Handle the case state==None when computing expt values
    rho0 = state if state else qutip.steadystate(H, c_ops)
    if is_e_op_hermitian:
        # Analitycal solution for x,x as operators.
        base = 0
        base += qutip.expect(a*x, rho0)*np.exp(-1j*w*times - gamma*times/2)
        base += qutip.expect(a.dag()*x, rho0)*np.exp(1j*w*times - gamma*times/2)
        base /= np.sqrt(2)
    else:
        # Analitycal solution for a,adag as operators.
        base = qutip.expect(a*a.dag(), rho0)*np.exp(-1j*w*times - gamma*times/2)

    cmp = qutip.correlation_2op_1t(H, state, times, c_ops, a_op, b_op, solver=solver)

    np.testing.assert_allclose(base, cmp, atol=0.25 if solver == 'mc' else 2e-5)


def test_correlation_timedependant_op():
    num = qutip.num(2)
    a = qutip.destroy(2)
    sx = qutip.sigmax()
    sz = qutip.sigmaz()
    times = np.arange(4)
    # switch between sx and sz at t=1.5
    A_op = qutip.QobjEvo([[sx, lambda t: t<=1.5], [sz, lambda t: t>1.5]])

    cmp_sx = qutip.correlation_2op_1t(num, None, times, [a], sx, sx)
    cmp_sz = qutip.correlation_2op_1t(num, None, times, [a], sz, sx)
    cmp_switch = qutip.correlation_2op_1t(num, None, times, [a], A_op, sx)
    np.testing.assert_allclose(cmp_sx[:2], cmp_switch[:2])
    np.testing.assert_allclose(cmp_sz[-2:], cmp_switch[-2:])


def test_alternative_solver():
    from qutip.solver.mesolve import MESolver
    from qutip.solver.brmesolve import BRSolver

    H = qutip.num(5)
    a = qutip.destroy(5)
    a_ops = [(a+a.dag(), qutip.coefficient(lambda _, w: w>0, args={"w":0}))]

    br = BRSolver(H, a_ops)
    me = MESolver(H, [a])
    times = np.arange(4)

    br_corr = qutip.correlation_3op(br, qutip.basis(5), [0], times, a, a.dag())
    me_corr = qutip.correlation_3op(me, qutip.basis(5), [0], times, a, a.dag())

    np.testing.assert_allclose(br_corr, me_corr)


def test_G1():
    H = qutip.Qobj([[0,1], [1,0]])
    psi0 = qutip.basis(2)
    taus = np.linspace(0, 1, 11)
    scale = 2
    a_op = qutip.sigmaz() * scale
    g1, G1 = qutip.coherence_function_g1(H, psi0, taus, [], a_op)
    expected = np.array([np.cos(t)**2 - np.sin(t)**2 for t in taus])
    np.testing.assert_allclose(g1, expected, rtol=2e-5)
    np.testing.assert_allclose(G1, expected * scale**2, rtol=2e-5)


def test_G2():
    N = 10
    H = qutip.rand_dm(N)
    psi0 = qutip.rand_ket(N)
    taus = np.linspace(0, 1, 11)
    scale = 2
    a_op = qutip.rand_unitary(N) * scale
    g1, G1 = qutip.coherence_function_g2(H, psi0, taus, [], a_op)
    expected = np.ones(11)
    np.testing.assert_allclose(g1, expected, rtol=2e-5)
    np.testing.assert_allclose(G1, expected * scale**4, rtol=2e-5)