# This file is part of QuTiP: Quantum Toolbox in Python.
#
#    Copyright (c) 2011 and later, The QuTiP Project
#    All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without
#    modification, are permitted provided that the following conditions are
#    met:
#
#    1. Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
#       of its contributors may be used to endorse or promote products derived
#       from this software without specific prior written permission.
#
#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###############################################################################
import itertools

import numpy as np
from numpy.testing import assert_
import pytest

from qutip import *
from qutip.legacy.ptrace import _ptrace as _pt


@pytest.fixture(params=[True, False], ids=['sparse', 'dense'])
def sparse(request):
    return request.param


@pytest.fixture(params=[True, False], ids=['dm', 'ket'])
def dm(request):
    return request.param


@pytest.fixture
def state(dm):
    dims = [2, 3, 4]
    state = qutip.rand_ket(np.prod(dims), dims=[dims, [1]*len(dims)])
    if dm:
        state = state.proj()
    return state


def test_ptrace_noncompound_rand(sparse, dm):
    """Test `A.ptrace(0) == A` when `A` is in a non-tensored Hilbert space."""
    for _ in range(5):
        state = qutip.rand_ket(5)
        if dm:
            state = state.proj()
        assert state.ptrace(0, sparse=sparse) == state


@pytest.mark.parametrize('pair', list(itertools.combinations(range(3), 2)))
def test_ptrace_unsorted_selection_subset(state, sparse, pair):
    """
    Regression test for gh-1325.  ptrace should work the same independently of
    the order of the input; no transposition in done in the trace operation.
    """
    # pair is always sorted.
    state_ordered = state.ptrace(pair, sparse=sparse)
    state_reversed = state.ptrace(pair[::-1], sparse=sparse)
    assert state_ordered.dims == state_reversed.dims
    assert state_ordered == state_reversed


@pytest.mark.parametrize('permutation', list(itertools.permutations(range(3))))
def test_ptrace_unsorted_selection_all(state, sparse, permutation):
    state_ptraced = state.ptrace(permutation, sparse=sparse)
    assert state.dims == state_ptraced.dims
    assert state == state_ptraced


@pytest.mark.parametrize(['selection', 'exception'], [
    pytest.param(4, IndexError, id='too big'),
    pytest.param(-1, IndexError, id='too small'),
    pytest.param([0, 0], ValueError, id='duplicate'),
    # 'too many' may throw either from duplication or invalid index.
    pytest.param([0, 1, 2, 3], Exception, id='too many'),
])
def test_ptrace_fails_on_invalid_input(state, sparse, selection, exception):
    with pytest.raises(exception):
        state.ptrace(selection, sparse=sparse)


def test_ptrace_rand():
    'ptrace : randomized tests, sparse'
    for k in range(5):
        A = tensor(rand_ket(5), rand_ket(2), rand_ket(3))
        B = A.ptrace([1,2], True)
        bdat,bd,bs = _pt(A, [1,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,2], True)
        bdat,bd,bs = _pt(A, [0,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,1], True)
        bdat,bd,bs = _pt(A, [0,1])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

    for k in range(5):
        A = tensor(rand_dm(2), thermal_dm(10,1), rand_unitary(3))
        B = A.ptrace([1,2], True)
        bdat,bd,bs = _pt(A, [1,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,2], True)
        bdat,bd,bs = _pt(A, [0,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,1], True)
        bdat,bd,bs = _pt(A, [0,1])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

    for k in range(5):
        A = tensor(rand_ket(2),rand_ket(2),rand_ket(2),
                    rand_ket(2),rand_ket(2),rand_ket(2))
        B = A.ptrace([3,2], True)
        bdat,bd,bs = _pt(A, [3,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,2], True)
        bdat,bd,bs = _pt(A, [0,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,1], True)
        bdat,bd,bs = _pt(A, [0,1])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

    for k in range(5):
        A = rand_dm(64,0.5,dims=[[4,4,4],[4,4,4]])
        B = A.ptrace([0], True)
        bdat,bd,bs = _pt(A, [0])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([1], True)
        bdat,bd,bs = _pt(A, [1])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,2], True)
        bdat,bd,bs = _pt(A, [0,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)


def test_ptrace_rand():
    'ptrace : randomized tests, dense'
    for k in range(5):
        A = tensor(rand_ket(5), rand_ket(2), rand_ket(3))
        B = A.ptrace([1,2], False)
        bdat,bd,bs = _pt(A, [1,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,2], False)
        bdat,bd,bs = _pt(A, [0,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,1], False)
        bdat,bd,bs = _pt(A, [0,1])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

    for k in range(5):
        A = tensor(rand_dm(2), thermal_dm(10,1), rand_unitary(3))
        B = A.ptrace([1,2], False)
        bdat,bd,bs = _pt(A, [1,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,2], False)
        bdat,bd,bs = _pt(A, [0,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,1], False)
        bdat,bd,bs = _pt(A, [0,1])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

    for k in range(5):
        A = tensor(rand_ket(2),rand_ket(2),rand_ket(2),
                    rand_ket(2),rand_ket(2),rand_ket(2))
        B = A.ptrace([3,2], False)
        bdat,bd,bs = _pt(A, [3,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,2], False)
        bdat,bd,bs = _pt(A, [0,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,1], False)
        bdat,bd,bs = _pt(A, [0,1])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

    for k in range(5):
        A = rand_dm(64,0.5,dims=[[4,4,4],[4,4,4]])
        B = A.ptrace([0], False)
        bdat,bd,bs = _pt(A, [0])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([1], False)
        bdat,bd,bs = _pt(A, [1])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,2], False)
        bdat,bd,bs = _pt(A, [0,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)
