import pytest
import itertools
import random
import numpy as np
import qutip
import scipy.sparse

from qutip.core.dimensions import flatten
from qutip.core.energy_restricted import EnrSpace, Space


def _n_enr_states(dimensions, n_excitations):
    """
    Calculate the total number of distinct ENR states for a given set of
    subspaces.  This method is not intended to be fast or efficient, it's
    intended to be obviously correct for testing purposes.
    """
    count = 0
    for excitations in itertools.product(*map(range, dimensions)):
        count += int(sum(excitations) <= n_excitations)
    return count


@pytest.fixture(params=[
    pytest.param([4], id="single"),
    pytest.param([4]*2, id="tensor-equal-2"),
    pytest.param([4]*3, id="tensor-equal-3"),
    pytest.param([4]*4, id="tensor-equal-4"),
    pytest.param([2, 3, 4], id="tensor-not-equal"),
])
def dimensions(request):
    return request.param


@pytest.fixture(params=[1, 2, 3, 1_000_000])
def n_excitations(request):
    return request.param


class TestOperator:
    def test_no_restrictions(self, dimensions):
        """
        Test that the restricted-excitation operators are equal to the standard
        operators when there aren't any restrictions.
        """
        test_operators = qutip.enr_destroy(dimensions, sum(dimensions))
        a = [qutip.destroy(n) for n in dimensions]
        iden = [qutip.qeye(n) for n in dimensions]
        for i, test in enumerate(test_operators):
            expected = qutip.tensor(iden[:i] + [a[i]] + iden[i+1:])
            assert test.data == expected.data
            assert test.dims == [dimensions, dimensions]

    def test_space_size_reduction(self, dimensions, n_excitations):
        test_operators = qutip.enr_destroy(dimensions, n_excitations)
        expected_size = _n_enr_states(dimensions, n_excitations)
        expected_shape = (expected_size, expected_size)
        for test in test_operators:
            assert test.shape == expected_shape
            assert test.dims == [dimensions, dimensions]

    def test_identity(self, dimensions, n_excitations):
        iden = qutip.enr_identity(dimensions, n_excitations)
        expected_size = _n_enr_states(dimensions, n_excitations)
        expected_shape = (expected_size, expected_size)
        assert np.all(iden.diag() == 1)
        assert np.all(iden.full() - np.diag(iden.diag()) == 0)
        assert iden.shape == expected_shape
        assert iden.dims == [dimensions, dimensions]

    def test_nstates(self, dimensions, n_excitations):
        expected_size = _n_enr_states(dimensions, n_excitations)
        nstates = qutip.enr_nstates(dimensions, n_excitations)
        assert nstates == expected_size


class TestEnrSpace:
    @pytest.mark.parametrize(["base", "flat"], [
        pytest.param(EnrSpace([2, 2], 1), [3], id="single"),
        pytest.param(Space(Space(2), Space(EnrSpace([2, 2], 1), Space(4))), [2, 3, 4], id="nested"),
    ])
    def test_flatten_collapses(self, base, flat):
        assert flatten(base) == flat
        assert base.flat() == flat

    @pytest.mark.parametrize(["space", "expected"], [
        pytest.param(EnrSpace([2, 2], 1), EnrSpace([2, 2], 1), id="no-scalars"),
        pytest.param(EnrSpace([2, 1, 3], 2), EnrSpace([2, 3], 2), id="simple"),
        pytest.param(Space(Space(5), Space(1), EnrSpace([2, 3, 1], 1)),
                     Space(Space(5), EnrSpace([2, 3], 1)),
                     id="nested")
    ])
    def test_drop_scalar_dims(self, space, expected):
        assert space.drop_scalar_dims() == expected


def test_fock_state(dimensions, n_excitations):
    """
    Test Fock state creation agrees with the number operators implied by the
    existence of the ENR annihiliation operators.
    """
    number = [a.dag()*a for a in qutip.enr_destroy(dimensions, n_excitations)]
    bases = list(qutip.state_number_enumerate(dimensions, n_excitations))
    n_samples = min((len(bases), 5))
    for basis in random.sample(bases, n_samples):
        state = qutip.enr_fock(dimensions, n_excitations, basis)
        for n, x in zip(number, basis):
            assert abs(n.matrix_element(state.dag(), state)) - x < 1e-10


def test_fock_state_error():
    with pytest.raises(ValueError) as e:
        state = qutip.enr_fock([2, 2, 2], 1, [1, 1, 1])
    assert str(e.value).startswith("state tuple ")


def _reference_dm(dimensions, n_excitations, nbars):
    """
    Get the reference density matrix explicitly, to compare to the direct ENR
    construction.
    """
    if np.isscalar(nbars):
        nbars = [nbars] * len(dimensions)
    keep = np.array([
        i for i, state in enumerate(qutip.state_number_enumerate(dimensions))
        if sum(state) <= n_excitations
    ])
    dm = qutip.tensor([qutip.thermal_dm(dimension, nbar)
                       for dimension, nbar in zip(dimensions, nbars)])
    out = qutip.Qobj(dm.full()[keep[:, None], keep[None, :]])
    return out / out.tr()


@pytest.mark.parametrize("nbar_type", ["scalar", "vector"])
def test_thermal_dm(dimensions, n_excitations, nbar_type):
    # Ensure that the average number of excitations over all the states is
    # much less than the total number of allowed excitations.
    if nbar_type == "scalar":
        nbars = 0.1 * n_excitations / len(dimensions)
    else:
        nbars = np.random.rand(len(dimensions))
        nbars *= (0.1 * n_excitations) / np.sum(nbars)
    test_dm = qutip.enr_thermal_dm(dimensions, n_excitations, nbars)
    expect_dm = _reference_dm(dimensions, n_excitations, nbars)
    np.testing.assert_allclose(test_dm.full(), expect_dm.full(), atol=1e-12)


def test_mesolve_ENR():
    # Ensure ENR states work with mesolve
    # We compare the output to an exact truncation of the
    # single-excitation Jaynes-Cummings model
    eps = 2 * np.pi
    omega_c = 2 * np.pi
    g = 0.1 * omega_c
    gam = 0.01 * omega_c
    tlist = np.linspace(0, 20, 100)
    N_cut = 2

    sz = qutip.sigmaz() & qutip.qeye(N_cut)
    sm = qutip.destroy(2).dag() & qutip.qeye(N_cut)
    a = qutip.qeye(2) & qutip.destroy(N_cut)
    H_JC = (0.5 * eps * sz + omega_c * a.dag()*a +
            g * (a * sm.dag() + a.dag() * sm))
    psi0 = qutip.basis(2, 0) & qutip.basis(N_cut, 0)
    c_ops = [np.sqrt(gam) * a]

    result_JC = qutip.mesolve(H_JC, psi0, tlist, c_ops, e_ops=[sz])

    N_exc = 1
    dims = [2, N_cut]
    d = qutip.enr_destroy(dims, N_exc)
    sz = 2*d[0].dag()*d[0]-1
    b = d[0]
    a = d[1]
    psi0 = qutip.enr_fock(dims, N_exc, [1, 0])
    H_enr = (eps * b.dag()*b + omega_c * a.dag() * a +
             g * (b.dag() * a + a.dag() * b))
    c_ops = [np.sqrt(gam) * a]

    result_enr = qutip.mesolve(H_enr, psi0, tlist, c_ops, e_ops=[sz])

    np.testing.assert_allclose(result_JC.expect[0],
                               result_enr.expect[0], atol=1e-2)


def test_steadystate_ENR():
    # Ensure ENR states work with steadystate functions
    # We compare the output to an exact truncation of the
    # single-excitation Jaynes-Cummings model
    eps = 2 * np.pi
    omega_c = 2 * np.pi
    g = 0.1 * omega_c
    gam = 0.01 * omega_c
    N_cut = 2

    sz = qutip.sigmaz() & qutip.qeye(N_cut)
    sm = qutip.destroy(2).dag() & qutip.qeye(N_cut)
    a = qutip.qeye(2) & qutip.destroy(N_cut)
    H_JC = (0.5 * eps * sz + omega_c * a.dag()*a +
            g * (a * sm.dag() + a.dag() * sm))
    c_ops = [np.sqrt(gam) * a]

    result_JC = qutip.steadystate(H_JC, c_ops)
    exp_sz_JC = qutip.expect(sz, result_JC)

    N_exc = 1
    dims = [2, N_cut]
    d = qutip.enr_destroy(dims, N_exc)
    sz = 2*d[0].dag()*d[0]-1
    b = d[0]
    a = d[1]
    H_enr = (eps * b.dag()*b + omega_c * a.dag() * a +
             g * (b.dag() * a + a.dag() * b))
    c_ops = [np.sqrt(gam) * a]

    result_enr = qutip.steadystate(H_enr, c_ops)
    exp_sz_enr = qutip.expect(sz, result_enr)

    np.testing.assert_allclose(exp_sz_JC,
                               exp_sz_enr, atol=1e-2)


def test_eigenstates_ENR():
    a1, a2 = qutip.enr_destroy([2, 2], 1)
    H = a1.dag() * a1 + a2.dag() * a2
    vals, vecs = H.eigenstates()
    for val, vec in zip(vals, vecs):
        assert val * vec == H @ vec


# function to convert an ENR Qobj to a regular Qobj
def enr_to_reg(q):
    if q.type == "scalar":
        return q
    dims = q.dims[0]
    n_ex = q._dims[0].n_excitations

    if q.isoper:
        newdims = [list(dims)]*2
        newshape = [np.prod(dims)]*2
    elif q.isket:
        newdims = [list(dims), [1]*len(dims)]
        newshape = [np.prod(dims), 1]
    else:
        raise ValueError("Invalid Qobj type")

    newindices = [i for i, state  # get indices for unrestricted space
                  in enumerate(qutip.state_number_enumerate(dims))
                  if sum(state) <= n_ex]
    dtype = q.dtype
    if dtype is qutip.data.Dense:
        mat = q.data_as()
        newmat = np.zeros(newshape, dtype=mat.dtype)
        if q.isoper:
            newmat[np.ix_(newindices, newindices)] = mat
        else:
            newmat[newindices] = mat
    else:
        mat = q.to("csr").data_as().tocoo()
        newrows = [newindices[i] for i in mat.row]
        newcols = [newindices[i] for i in mat.col]
        newmat = scipy.sparse.coo_matrix(
            (mat.data, (newrows, newcols)), shape=newshape)

    return qutip.Qobj(newmat, dims=newdims, dtype=dtype)


@pytest.fixture(params=["csr", "dense", "dia"])
def dtype(request):
    return request.param


# copied parameters from test_ptrace.py
@pytest.mark.parametrize('dims, sel',
                         [
                             ([5, 2, 3], [2, 1]),
                             ([5, 2, 3], [0, 2]),
                             ([5, 2, 3], [0, 1]),
                             ([2]*6, [3, 2]),
                             ([2]*6, [0, 2]),
                             ([2]*6, [0, 1]),
                             ([4, 4, 4], []),
                             ([4, 4, 4], [0, 1, 2]),
                         ])
@pytest.mark.filterwarnings("ignore:enr_ptrace")  
def test_enr_ptrace(dims, sel, n_excitations, dtype):
    nstates_enr = _n_enr_states(dims, n_excitations)
    # use qutip to make a random sparse Hermitian matrix w/ trace 1
    rho_in_enr = qutip.rand_dm(
        nstates_enr, 0.05, distribution="herm", dtype=dtype)

    rho_in_enr.dims = [qutip.energy_restricted.EnrSpace(dims, n_excitations)]*2
    rho_in_reg = enr_to_reg(rho_in_enr)

    rho_out_enr = qutip.enr_ptrace(rho_in_enr, sel)
    rho_out_toreg = enr_to_reg(rho_out_enr)
    rho_out_check = qutip.ptrace(rho_in_reg, sel)

    assert rho_out_toreg == rho_out_check


@pytest.fixture(params=[
    pytest.param(0, id="ket"),
    pytest.param(1, id="oper"),
])
def isoper(request):
    return request.param


@pytest.mark.parametrize('dims_list, n_ex_list',
                         [
                             ([[3]], [2]),
                             ([[3], [4]], [2, 1]),
                             ([[3, 2], [2, 2]], [2, 1]),
                             ([[3, 2], [4], [2, 2]], [2, 3, 2]),
                         ])
def test_enr_tensor(dims_list, n_ex_list, isoper, dtype):
    # isoper = 1 to test for operators, 0 to test for kets
    # figure out how big the matrices/vectors will be
    nstates_list = [_n_enr_states(dims, n_ex)
                    for dims, n_ex in zip(dims_list, n_ex_list)]
    # make sure sparse matrices are dense enough to not be empty
    dens_list = [max(0.1, 1/(nstates**(isoper+1))) for nstates in nstates_list]
    # number of columns in the random matrices, depending on if operator or ket
    ncol_arr = nstates_list if isoper else [1]*len(nstates_list)

    # generate the random matrices
    rand_mat_list = [scipy.sparse.random(
        nstates, ncol,
        density=dens, dtype="complex128")
        for (nstates, ncol, dens)
        in zip(nstates_list, ncol_arr, dens_list)]

    # figure out the row and column spaces for the ENR Qobj's
    rowspace_list = [qutip.energy_restricted.EnrSpace(dims, n_ex)
                     for dims, n_ex in zip(dims_list, n_ex_list)]
    if isoper:
        colspace_list = rowspace_list
    else:
        colspace_list = [[1]] * len(rowspace_list)

    # turn the random matrices into Qobj's
    q_in_enr_list = [
        qutip.Qobj(rand_mat, dims=[rowspace, colspace], dtype=dtype)
        for rand_mat, rowspace, colspace
        in zip(rand_mat_list, rowspace_list, colspace_list)]
    assert all(q_in_enr.isoper == isoper for q_in_enr in q_in_enr_list)

    q_in_reg_list = [enr_to_reg(q_in_enr) for q_in_enr in q_in_enr_list]

    q_out_enr = qutip.enr_tensor(*q_in_enr_list, newexcitations=sum(n_ex_list))
    q_out_toreg = enr_to_reg(q_out_enr)
    q_out_check = qutip.tensor(*q_in_reg_list)

    assert q_out_toreg == q_out_check


@pytest.mark.parametrize('truncate', [True, False])
def test_enr_tensor_trunc(isoper, dtype, truncate):
    q_in = qutip.enr_fock([2], 1, [1], dtype=dtype)
    if isoper:
        q_in = q_in.proj()

    if not truncate:
        with pytest.raises(ValueError) as e:
            qutip.enr_tensor(q_in, q_in)
        assert "not in the new restricted state space" in str(e.value)
    else:
        q_out = qutip.enr_tensor(q_in, q_in, truncate=truncate, verbose=False)
        q_out_check = qutip.zero_ket(
            qutip.energy_restricted.EnrSpace([2, 2], 1), dtype=dtype)
        if isoper:
            q_out_check = q_out_check.proj()
        assert q_out == q_out_check


def test_enr_coefficient_mul():
    # Ensure that QobjEvo can be multiplied by Coefficient with ENR operators
    a, _ = qutip.enr_destroy([3, 3], 2)
    evo = qutip.QobjEvo([a.dag() * a, qutip.coefficient("sin(2*pi*t)")])
    evo *= qutip.coefficient("t")
    assert evo.num_elements == 1
    assert evo(0.25) == 0.25 * a.dag() * a
