from . import test_mathematics as testing
import numpy as np
import scipy.linalg
import pytest
from qutip import data
from qutip.core.data import CSR, Dense, Dia


class TestPtrace(testing.UnaryOpMixin):
    def op_numpy(self, matrix, dims, sel):
        sel.sort()
        ndims = len(dims)
        dkeep = [dims[x] for x in sel]
        qtrace = list(set(range(ndims)) - set(sel))
        dtrace = [dims[x] for x in qtrace]

        matrix = matrix.reshape(dims + dims)
        matrix = matrix.transpose(
            qtrace + [ndims + i for i in qtrace] + sel + [ndims + i for i in sel]
        )
        matrix = matrix.reshape(
            [
                np.prod(dtrace, dtype=int),
                np.prod(dtrace, dtype=int),
                np.prod(dkeep, dtype=int),
                np.prod(dkeep, dtype=int),
            ]
        )
        return np.trace(matrix)

    # Custom shapes to have also custom dims and sel arguments.
    # These values should not be changed.
    dims = [2] * 7
    shapes = [(pytest.param((np.prod(dims), np.prod(dims))),)]
    bad_shapes = testing.shapes_not_square(np.prod(dims))

    specialisations = [
        pytest.param(data.ptrace_csr, CSR, CSR),
        pytest.param(data.ptrace_csr_dense, CSR, Dense),
        pytest.param(data.ptrace_dense, Dense, Dense),
        pytest.param(data.ptrace_dia, Dia, Dia),
    ]

    @pytest.mark.parametrize(
        "sel",
        [[0], [0, 3, 6], [0, 6, 3], list(range(len(dims))), []],
        ids=[
            "keep_one",
            "keep_multiple_sorted",
            "keep_multiple_unsorted",
            "trace_none",
            "trace_all",
        ],
    )
    def test_mathematically_correct(self, op, data_m, out_type, sel):
        """
        Test that the unary operation is mathematically correct for all the
        known type specialisations.
        """
        matrix = data_m()
        expected = self.op_numpy(matrix.to_array(), self.dims, sel)
        test = op(matrix, self.dims, sel)
        assert isinstance(test, out_type)
        assert test.shape == expected.shape
        np.testing.assert_allclose(
            test.to_array(), expected, atol=self.atol, rtol=self.rtol
        )

    def test_incorrect_shape_raises(self, op, data_m):
        """
        Test that the operation produces a suitable error if the shape of the
        operand is not square.
        """
        with pytest.raises(ValueError):
            op(data_m(), self.dims, sel=[0, 1])

    # `out_type` is included but not used so that
    # `generate_mathematically_correct` can be re-used.
    @pytest.mark.parametrize(
        "dims",
        [[2], [0], [-2, -2] + [2] * 5, [1.2, 2.2, 3.3]],
        ids=[
            "dims_different_to_shape",
            "dims_0",
            "dims_prod_is_shape_but_negative",
            "dims_is_not_int",
        ],
    )
    def test_incorrect_dims_raises(self, op, data_m, out_type, dims):
        with pytest.raises(ValueError):
            op(data_m(), dims, sel=[0, 1])

    def generate_incorrect_dims_raises(self, metafunc):
        self.generate_mathematically_correct(metafunc)

    @pytest.mark.parametrize(
        "sel",
        [[2, 10], [-1, 2]],
        ids=[
            "sel_value_larger_than_dims",
            "sel_value_negative",
        ],
    )
    def test_incorrect_sel_raises(self, op, data_m, out_type, sel):
        with pytest.raises(IndexError):
            op(data_m(), dims=self.dims, sel=sel)

    def generate_incorrect_sel_raises(self, metafunc):
        self.generate_mathematically_correct(metafunc)