import pytest
import numpy as np
from qutip import (
    mesolve, liouvillian, QobjEvo, spre, spost, CoreOptions,
    destroy, coherent, qeye, fock_dm, num, basis
)
from qutip.solver.stochastic import smesolve, ssesolve, SMESolver, SSESolver
from qutip.core import data as _data


def f(t, w):
    return w * t

def _make_system(N, system):
    gamma = 0.25
    a = destroy(N)

    if system == "no sc_ops":
        H = a.dag() * a
        sc_ops = []

    if system == "simple":
        H = a.dag() * a
        sc_ops = [np.sqrt(gamma) * a]

    elif system == "2 c_ops":
        H = QobjEvo([a.dag() * a])
        sc_ops = [np.sqrt(gamma) * a, gamma * a * a]

    elif system == "H td":
        H = [[a.dag() * a, f]]
        sc_ops = [np.sqrt(gamma) * QobjEvo(a)]

    elif system == "complex":
        H = [a.dag() * a + a.dag() + a]
        sc_ops = [np.sqrt(gamma) * a, gamma * a * a]

    elif system == "c_ops td":
        H = [a.dag() * a]
        sc_ops = [[np.sqrt(gamma) * a, f]]

    return H, sc_ops


@pytest.mark.parametrize("system", [
    "no sc_ops", "simple", "2 c_ops", "H td", "complex", "c_ops td",
])
@pytest.mark.parametrize("heterodyne", [True, False])
def test_smesolve(heterodyne, system):
    tol = 0.05
    N = 4
    ntraj = 20

    H, sc_ops = _make_system(N, system)
    c_ops = [destroy(N)]
    psi0 = coherent(N, 0.5)
    a = destroy(N)
    e_ops = [a.dag() * a, a + a.dag(), (-1j)*(a - a.dag())]

    times = np.linspace(0, 0.1, 21)
    res_ref = mesolve(H, psi0, times, c_ops + sc_ops,
                      e_ops=e_ops, args={"w": 2})

    options = {
        "store_measurement": False,
        "map": "serial",
    }

    res = smesolve(
        H, psi0, times, sc_ops=sc_ops, e_ops=e_ops, c_ops=c_ops,
        ntraj=ntraj, args={"w": 2}, options=options, heterodyne=heterodyne,
        seeds=1,
    )

    for idx in range(len(e_ops)):
        np.testing.assert_allclose(
            res.expect[idx], res_ref.expect[idx], rtol=tol, atol=tol
        )


@pytest.mark.parametrize("system", [
    "no sc_ops", "simple"
])
@pytest.mark.parametrize("heterodyne", [True, False])
@pytest.mark.parametrize("method", SMESolver.avail_integrators().keys())
def test_smesolve_methods(method, heterodyne, system):
    tol = 0.05
    N = 4
    ntraj = 20

    H, sc_ops = _make_system(N, system)
    c_ops = [destroy(N)]
    psi0 = coherent(N, 0.5)
    a = destroy(N)
    e_ops = [a.dag() * a, a + a.dag(), (-1j)*(a - a.dag())]

    times = np.linspace(0, 0.1, 21)
    res_ref = mesolve(H, psi0, times, c_ops + sc_ops,
                      e_ops=e_ops, args={"w": 2})

    options = {
        "store_measurement": True,
        "map": "parallel",
        "method": method,
    }

    res = smesolve(
        H, psi0, times, sc_ops=sc_ops, e_ops=e_ops, c_ops=c_ops,
        ntraj=ntraj, args={"w": 2}, options=options, heterodyne=heterodyne,
        seeds=list(range(ntraj)),
    )

    for idx in range(len(e_ops)):
        np.testing.assert_allclose(
            res.expect[idx], res_ref.expect[idx], rtol=tol, atol=tol
        )

    assert len(res.measurement) == ntraj

    if heterodyne:
        assert all([
            dw.shape == (len(sc_ops), 2, len(times)-1)
            for dw in res.dW
        ])
        assert all([
            w.shape == (len(sc_ops), 2, len(times))
            for w in res.wiener_process
        ])
        assert all([
            m.shape == (len(sc_ops), 2, len(times)-1)
            for m in res.measurement
        ])
    else:
        assert all([
            dw.shape == (len(sc_ops), len(times)-1)
            for dw in res.dW
        ])
        assert all([
            w.shape == (len(sc_ops), len(times))
            for w in res.wiener_process
        ])
        assert all([
            m.shape == (len(sc_ops), len(times)-1)
            for m in res.measurement
        ])


@pytest.mark.parametrize("system", [
    "no sc_ops", "simple", "2 c_ops", "H td", "complex", "c_ops td",
])
@pytest.mark.parametrize("heterodyne", [True, False])
def test_ssesolve(heterodyne, system):
    tol = 0.1
    N = 4
    ntraj = 20

    H, sc_ops = _make_system(N, system)
    psi0 = coherent(N, 0.5)
    a = destroy(N)
    e_ops = [a.dag() * a, a + a.dag(), (-1j)*(a - a.dag())]

    times = np.linspace(0, 0.1, 21)
    res_ref = mesolve(H, psi0, times, sc_ops, e_ops=e_ops, args={"w": 2})

    options = {
        "map": "serial",
    }

    res = ssesolve(
        H, psi0, times, sc_ops, e_ops=e_ops,
        ntraj=ntraj, args={"w": 2}, options=options, heterodyne=heterodyne,
        seeds=list(range(ntraj)),
    )

    for idx in range(len(e_ops)):
        np.testing.assert_allclose(
            res.expect[idx], res_ref.expect[idx], rtol=tol, atol=tol
        )

    assert res.measurement is None
    assert res.wiener_process is None
    assert res.dW is None


@pytest.mark.parametrize("system", [
    "no sc_ops", "simple"
])
@pytest.mark.parametrize("heterodyne", [True, False])
@pytest.mark.parametrize("method", SSESolver.avail_integrators().keys())
def test_ssesolve_method(method, heterodyne, system):
    "Stochastic: smesolve: homodyne, time-dependent H"
    tol = 0.1
    N = 4
    ntraj = 20

    H, sc_ops = _make_system(N, system)
    psi0 = coherent(N, 0.5)
    a = destroy(N)
    e_ops = [a.dag() * a, a + a.dag(), (-1j)*(a - a.dag())]

    times = np.linspace(0, 0.1, 21)
    res_ref = mesolve(H, psi0, times, sc_ops, e_ops=e_ops, args={"w": 2})

    options = {
        "store_measurement": True,
        "map": "parallel",
        "method": method,
        "keep_runs_results": True,
    }

    res = ssesolve(
        H, psi0, times, sc_ops, e_ops=e_ops,
        ntraj=ntraj, args={"w": 2}, options=options, heterodyne=heterodyne,
        seeds=1,
    )

    for idx in range(len(e_ops)):
        np.testing.assert_allclose(
            res.average_expect[idx], res_ref.expect[idx], rtol=tol, atol=tol
        )

    assert len(res.measurement) == ntraj

    if heterodyne:
        assert all([
            dw.shape == (len(sc_ops), 2, len(times)-1)
            for dw in res.dW
        ])
        assert all([
            w.shape == (len(sc_ops), 2, len(times))
            for w in res.wiener_process
        ])
        assert all([
            m.shape == (len(sc_ops), 2, len(times)-1)
            for m in res.measurement
        ])
    else:
        assert all([
            dw.shape == (len(sc_ops), len(times)-1)
            for dw in res.dW
        ])
        assert all([
            w.shape == (len(sc_ops), len(times))
            for w in res.wiener_process
        ])
        assert all([
            m.shape == (len(sc_ops), len(times)-1)
            for m in res.measurement
        ])


def test_reuse_seeds():
    tol = 0.05
    N = 4
    ntraj = 5

    H, sc_ops = _make_system(N, "simple")
    c_ops = [destroy(N)]
    psi0 = coherent(N, 0.5)
    a = destroy(N)
    e_ops = [a.dag() * a, a + a.dag(), (-1j)*(a - a.dag())]

    times = np.linspace(0, 0.1, 2)

    options = {
        "store_final_state": True,
        "map": "serial",
        "keep_runs_results": True,
        "store_measurement": True,
    }

    res = smesolve(
        H, psi0, times, sc_ops=sc_ops, e_ops=e_ops, c_ops=c_ops,
        ntraj=ntraj, args={"w": 2}, options=options,
    )

    res2 = smesolve(
        H, psi0, times, sc_ops=sc_ops, e_ops=e_ops, c_ops=c_ops,
        ntraj=ntraj, args={"w": 2}, options=options,
        seeds=res.seeds,
    )

    np.testing.assert_allclose(
        res.wiener_process, res2.wiener_process, atol=1e-14
    )

    np.testing.assert_allclose(res.expect, res2.expect, atol=1e-14)

    for out1, out2 in zip(res.final_state, res2.final_state):
        assert out1 == out2


@pytest.mark.parametrize("heterodyne", [True, False])
def test_measurements(heterodyne):
    N = 10
    ntraj = 1

    H = num(N)
    sc_ops = [destroy(N)]
    psi0 = basis(N, N-1)

    times = np.linspace(0, 1.0, 11)

    solver = SMESolver(H, sc_ops, heterodyne=heterodyne)

    solver.options["store_measurement"] = "start"
    res_start = solver.run(psi0, times, ntraj=ntraj, seeds=1)

    solver.options["store_measurement"] = "middle"
    res_middle = solver.run(psi0, times, ntraj=ntraj, seeds=1)

    solver.options["store_measurement"] = "end"
    res_end = solver.run(psi0, times, ntraj=ntraj, seeds=1)

    diff = np.sum(np.abs(res_end.measurement[0] - res_start.measurement[0]))
    assert diff > 0.1 # Each measurement should be different by ~dt
    np.testing.assert_allclose(
        res_middle.measurement[0] * 2,
        res_start.measurement[0] + res_end.measurement[0],
    )

    np.testing.assert_allclose(
        np.diff(res_start.wiener_process[0][0]), res_start.dW[0][0]
    )


@pytest.mark.parametrize("heterodyne", [True, False])
def test_m_ops(heterodyne):
    N = 10

    H = num(N)
    sc_ops = [destroy(N), qeye(N)]
    psi0 = basis(N, N-1)
    m_ops = [num(N), qeye(N)]
    if heterodyne:
        m_ops = m_ops * 2

    times = np.linspace(0, 1.0, 51)

    options = {"store_measurement": "end",}

    solver = SMESolver(H, sc_ops, heterodyne=heterodyne, options=options)
    solver.m_ops = m_ops
    solver.dW_factors = [0.] * len(m_ops)

    res = solver.run(psi0, times, e_ops=m_ops)
    # With dW_factors=0, measurements are computed as expectation values.
    if heterodyne:
        np.testing.assert_allclose(res.expect[0][1:], res.measurement[0][0][0])
        np.testing.assert_allclose(res.expect[1][1:], res.measurement[0][0][1])
    else:
        np.testing.assert_allclose(res.expect[0][1:], res.measurement[0][0])
        np.testing.assert_allclose(res.expect[1][1:], res.measurement[0][1])

    solver.dW_factors = [1.] * len(m_ops)
    # With dW_factors=0, measurements are computed as expectation values.
    res = solver.run(psi0, times, e_ops=m_ops)
    std = 1 / times[1]**0.5
    if heterodyne:
        noise = res.expect[0][1:] - res.measurement[0][0][0]
    else:
        noise = res.expect[0][1:] - res.measurement[0][0]
    assert np.mean(noise) == pytest.approx(0., abs=std / 50**0.5 * 5)
    assert np.std(noise) == pytest.approx(std, abs=std / 50**0.5 * 5)


def test_feedback():
    N = 10
    ntraj = 2

    def func(t, A, W):
        return (A - 6) * (A.real > 6.) * W(t)[0]

    H = num(10)
    sc_ops = [QobjEvo(
        [destroy(N), func],
        args={
            "A": SMESolver.ExpectFeedback(num(10)),
            "W": SMESolver.WienerFeedback()
        }
    )]
    psi0 = basis(N, N-3)

    times = np.linspace(0, 2, 101)
    options = {"map": "serial", "dt": 0.0005}

    solver = SMESolver(H, sc_ops=sc_ops, heterodyne=False, options=options)
    results = solver.run(psi0, times, e_ops=[num(N)], ntraj=ntraj)

    # If this was deterministic, it should never go under `6`.
    # We add a tolerance ~dt due to the stochatic part.
    assert np.all(results.expect[0] > 6. - 0.001)
    assert np.all(results.expect[0][-20:] < 6.7)


def test_deprecation_warnings():
    with pytest.warns(FutureWarning, match=r'map_func'):
        ssesolve(qeye(2), basis(2), [0, 0.01], [qeye(2)], map_func=None)

    with pytest.warns(FutureWarning, match=r'progress_bar'):
        ssesolve(qeye(2), basis(2), [0, 0.01], [qeye(2)], progress_bar=None)

    with pytest.warns(FutureWarning, match=r'nsubsteps'):
        ssesolve(qeye(2), basis(2), [0, 0.01], [qeye(2)], nsubsteps=None)

    with pytest.warns(FutureWarning, match=r'map_func'):
        ssesolve(qeye(2), basis(2), [0, 0.01], [qeye(2)], map_func=None)

    with pytest.warns(FutureWarning, match=r'store_all_expect'):
        ssesolve(qeye(2), basis(2), [0, 0.01], [qeye(2)], store_all_expect=1)

    with pytest.warns(FutureWarning, match=r'store_measurement'):
        ssesolve(qeye(2), basis(2), [0, 0.01], [qeye(2)], store_measurement=1)

    with pytest.raises(TypeError) as err:
        ssesolve(qeye(2), basis(2), [0, 0.01], [qeye(2)], m_ops=1)
    assert '"m_ops" and "dW_factors"' in str(err.value)


@pytest.mark.parametrize("method", ["euler", "rouchon"])
def test_small_step_warnings(method):
    with pytest.warns(RuntimeWarning, match=r'under minimum'):
        smesolve(
            qeye(2), basis(2), [0, 0.0000001], [qeye(2)],
            options={"method": method}
        )


@pytest.mark.parametrize("method", ["euler", "platen"])
@pytest.mark.parametrize("heterodyne", [True, False])
def test_run_from_experiment_close(method, heterodyne):
    N = 5

    H = num(N)
    a = destroy(N)
    sc_ops = [a, a @ a + (a @ a).dag()]
    psi0 = basis(N, N-1)
    tlist = np.linspace(0, 0.1, 501)
    options = {
        "store_measurement": "start",
        "dt": tlist[1],
        "store_states": True,
        "method": method,
    }
    with CoreOptions(auto_real_casting=False):
        solver = SSESolver(H, sc_ops, heterodyne, options=options)
        res_forward = solver.run(psi0, tlist, 1, e_ops=[H])
        res_backward = solver.run_from_experiment(
            psi0, tlist, res_forward.dW[0], e_ops=[H]
        )
        res_measure = solver.run_from_experiment(
            psi0, tlist, res_forward.measurement[0],
            e_ops=[H], measurement=True
        )

    np.testing.assert_allclose(
        res_backward.measurement, res_forward.measurement[0], atol=1e-10
    )
    np.testing.assert_allclose(
        res_measure.measurement, res_forward.measurement[0], atol=1e-10
    )

    np.testing.assert_allclose(res_backward.dW, res_forward.dW[0], atol=1e-10)
    np.testing.assert_allclose(res_measure.dW, res_forward.dW[0], atol=1e-10)

    np.testing.assert_allclose(
        res_backward.expect, res_forward.expect, atol=1e-10
    )
    np.testing.assert_allclose(
        res_measure.expect, res_forward.expect, atol=1e-10
    )


@pytest.mark.parametrize(
    "method", ["euler", "milstein", "platen", "pred_corr"]
)
@pytest.mark.parametrize("heterodyne", [True, False])
def test_run_from_experiment_open(method, heterodyne):
    N = 10

    H = num(N)
    a = destroy(N)
    sc_ops = [a, a.dag() * 0.1]
    psi0 = basis(N, N-1)
    tlist = np.linspace(0, 1, 251)
    options = {
        "store_measurement": "start",
        "dt": tlist[1],
        "store_states": True,
        "method": method,
    }
    solver = SMESolver(H, sc_ops, heterodyne, options=options)
    res_forward = solver.run(psi0, tlist, 1, e_ops=[H])
    res_backward = solver.run_from_experiment(
        psi0, tlist, res_forward.dW[0], e_ops=[H]
    )
    res_measure = solver.run_from_experiment(
        psi0, tlist, res_forward.measurement[0], e_ops=[H], measurement=True
    )

    np.testing.assert_allclose(
        res_backward.measurement, res_forward.measurement[0], atol=1e-10
    )
    np.testing.assert_allclose(
        res_measure.measurement, res_forward.measurement[0], atol=1e-10
    )

    np.testing.assert_allclose(res_backward.dW, res_forward.dW[0], atol=1e-10)
    np.testing.assert_allclose(res_measure.dW, res_forward.dW[0], atol=1e-10)

    np.testing.assert_allclose(
        res_backward.expect, res_forward.expect, atol=1e-10
    )
    np.testing.assert_allclose(
        res_measure.expect, res_forward.expect, atol=1e-10
    )


@pytest.mark.parametrize("store_measurement", [True, False])
@pytest.mark.parametrize("keep_runs_results", [True, False])
def test_merge_results(store_measurement, keep_runs_results):
    # Running smesolve with mixed ICs should be the same as running smesolve
    # multiple times and merging the results afterwards
    initial_state1 = basis([2, 2], [1, 0])
    initial_state2 = basis([2, 2], [1, 1])
    H = qeye([2, 2])
    L = destroy(2) & qeye(2)
    tlist = np.linspace(0, 1, 11)
    e_ops = [num(2) & qeye(2), qeye(2) & num(2)]

    options = {
        "store_measurement": True,
        "keep_runs_results": True,
        "store_states": True,
    }
    solver = SMESolver(H, [L], True, options=options)
    result1 = solver.run(initial_state1, tlist, 5, e_ops=e_ops)

    options = {
        "store_measurement": store_measurement,
        "keep_runs_results": keep_runs_results,
        "store_states": True,
    }
    solver = SMESolver(H, [L], True, options=options)
    result2 = solver.run(initial_state2, tlist, 10, e_ops=e_ops)

    result_merged = result1 + result2
    assert len(result_merged.seeds) == 15
    if store_measurement:
        with CoreOptions(atol=2e-8):
            # In the numpy 1.X test, this fail with the defautl atol=1e-12
            # One operation is computed in single precision?
            assert (
                result_merged.average_states[0] ==
                (initial_state1.proj() + 2 * initial_state2.proj()).unit()
            )
    np.testing.assert_allclose(result_merged.average_expect[0][0], 1)
    np.testing.assert_allclose(result_merged.average_expect[1], 2/3)

    if store_measurement:
        assert len(result_merged.measurement) == 15
        assert len(result_merged.dW) == 15
        assert all(
            dw.shape == result_merged.dW[0].shape
            for dw in result_merged.dW
        )
        assert len(result_merged.wiener_process) == 15
        assert all(
            w.shape == result_merged.wiener_process[0].shape
            for w in result_merged.wiener_process
        )


@pytest.mark.parametrize("open", [True, False])
@pytest.mark.parametrize("heterodyne", [True, False])
def test_step(open, heterodyne):
    state0 = basis(5, 3)
    kw = {}
    if open:
        SolverCls = SMESolver
        state0 = state0.proj()
    else:
        SolverCls = SSESolver

    solver = SolverCls(
        num(5),
        sc_ops=[destroy(5), destroy(5)**2 / 10],
        heterodyne=heterodyne,
        options={"dt": 0.001},
        **kw
    )
    solver.start(state0, t0=0)
    state1 = solver.step(0.01)
    assert state1.dims == state0.dims
    assert state1.norm() == pytest.approx(1, abs=0.01)
    state2, dW = solver.step(0.02, wiener_increment=True)
    assert state2.dims == state0.dims
    assert state2.norm() == pytest.approx(1, abs=0.01)
    if heterodyne:
        assert dW.shape == (2, 2)
        assert abs(dW[0, 0]) < 0.5 # 5 sigmas
    else:
        assert dW.shape == (2,)
        assert abs(dW[0]) < 0.5 # 5 sigmas
