# Source code for qutip.superop_reps

    # -*- coding: utf-8 -*-
    #
    # This module was initially contributed by Ben Criger.
    #
    """
    This module implements transformations between superoperator representations,
    including supermatrix, Kraus, Choi and Chi (process) matrix formalisms.
    """

    __all__ = ['super_to_choi', 'choi_to_super', 'choi_to_kraus', 'kraus_to_choi',
               'kraus_to_super', 'choi_to_chi', 'chi_to_choi', 'to_choi',
               'to_chi', 'to_super', 'to_kraus', 'to_stinespring'
               ]

    # Python Standard Library
    from itertools import starmap, product

    # NumPy/SciPy
    from numpy.core.multiarray import array, zeros
    from numpy.core.shape_base import hstack
    from numpy.matrixlib.defmatrix import matrix
    from numpy import sqrt, floor, log2
    from numpy import dot
    from scipy.linalg import eig, svd
    # Needed to avoid conflict with itertools.product.
    import numpy as np

    # Other QuTiP functions and classes
    from qutip.superoperator import vec2mat, operator_to_vector, sprepost
    from qutip.operators import identity, sigmax, sigmay, sigmaz
    from qutip.tensor import tensor, flatten
    from qutip.qobj import Qobj
    from qutip.states import basis

    # SPECIFIC SUPEROPERATORS -----------------------------------------------------

    def _dep_super(pe):
        """
        Returns the superoperator corresponding to qubit depolarization for a
        given parameter pe.

        TODO: if this is going into production (hopefully it isn't) then check
        CPTP, expand to arbitrary dimensional systems, etc.
        """
        return Qobj(dims=[[[2], [2]], [[2], [2]]],
                    inpt=array([[1. - pe / 2., 0., 0., pe / 2.],
                                [0., 1. - pe, 0., 0.],
                                [0., 0., 1. - pe, 0.],
                                [pe / 2., 0., 0., 1. - pe / 2.]]))

    def _dep_choi(pe):
        """
        Returns the choi matrix corresponding to qubit depolarization for a
        given parameter pe.

        TODO: if this is going into production (hopefully it isn't) then check
        CPTP, expand to arbitrary dimensional systems, etc.
        """
        return Qobj(dims=[[[2], [2]], [[2], [2]]],
                    inpt=array([[1. - pe / 2., 0., 0., 1. - pe],
                                [0., pe / 2., 0., 0.],
                                [0., 0., pe / 2., 0.],
                                [1. - pe, 0., 0., 1. - pe / 2.]]),
                    superrep='choi')

    # CHANGE OF BASIS FUNCTIONS ---------------------------------------------------
    # These functions find change of basis matrices, and are useful in converting
    # between (for instance) Choi and chi matrices. At some point, these should
    # probably be moved out to another module.

    _SINGLE_QUBIT_PAULI_BASIS = (identity(2), sigmax(), sigmay(), sigmaz())

    def _pauli_basis(nq=1):
        # NOTE: This is slow as can be.
        # TODO: Make this sparse. CSR format was causing problems for the [idx, :]
        #       slicing below.
        B = zeros((4 ** nq, 4 ** nq), dtype=complex)
        dims = [[[2] * nq] * 2] * 2

        for idx, op in enumerate(starmap(tensor,
                                         product(_SINGLE_QUBIT_PAULI_BASIS,
                                                 repeat=nq))):
            B[:, idx] = operator_to_vector(op).dag().full()

        return Qobj(B, dims=dims)

    # PRIVATE CONVERSION FUNCTIONS ------------------------------------------------
    # These functions handle the main work of converting between representations,
    # and are exposed below by other functions that add postconditions about types.
    #
    # TODO: handle type='kraus' as a three-index Qobj, rather than as a list?

    def _super_tofrom_choi(q_oper):
        """
        We exploit that the basis transformation between Choi and supermatrix
        representations squares to the identity, so that if we munge Qobj.type,
        we can use the same function.

        Since this function doesn't respect :attr:`Qobj.type`, we mark it as
        private; only those functions which wrap this in a way so as to preserve
        type should be called externally.
        """
        data = q_oper.data.toarray()
        dims = q_oper.dims
        new_dims = [[dims[1][1], dims[0][1]], [dims[1][0], dims[0][0]]]
        d0 = np.prod(np.ravel(new_dims[0]))
        d1 = np.prod(np.ravel(new_dims[1]))
        s0 = np.prod(dims[0][0])
        s1 = np.prod(dims[1][1])
        return Qobj(dims=new_dims,
                    inpt=data.reshape([s0, s1, s0, s1]).
                    transpose(3, 1, 2, 0).reshape((d0, d1)))

    def _isqubitdims(dims):
        """Checks whether all entries in a dims list are integer powers of 2.

        Parameters
        ----------
        dims : nested list of ints
            Dimensions to be checked.

        Returns
        -------
        isqubitdims : bool
            True if and only if every member of the flattened dims
            list is an integer power of 2.
        """
        return all([
            2**floor(log2(dim)) == dim
            for dim in flatten(dims)
        ])

    def _super_to_superpauli(q_oper):
        """
        Converts a superoperator in the column-stacking basis to
        the Pauli basis (assuming qubit dimensions).

        This is an internal function, as QuTiP does not currently have
        a way to mark that superoperators are represented in the Pauli
        basis as opposed to the column-stacking basis; a Pauli-basis
        ``type='super'`` would thus break other conversion functions.
        """
        # Ensure we start with a column-stacking-basis superoperator.
        sqobj = to_super(q_oper)
        if not _isqubitdims(sqobj.dims):
            raise ValueError("Pauli basis is only defined for qubits.")
        nq = int(log2(sqobj.shape[0]) / 2)
        B = _pauli_basis(nq) / sqrt(2**nq)
        # To do this, we have to hack a bit and force the dims to match,
        # since the _pauli_basis function makes different assumptions
        # about indices than we need here.
        B.dims = sqobj.dims
        return (B.dag() * sqobj * B)

    def super_to_choi(q_oper):
        # TODO: deprecate and make private in favor of to_choi,
        # which looks at Qobj.type to determine the right conversion function.
        """
        Takes a superoperator to a Choi matrix
        TODO: Sanitize input, incorporate as method on Qobj if type=='super'
        """
        q_oper = _super_tofrom_choi(q_oper)
        q_oper.superrep = 'choi'
        return q_oper

    def choi_to_super(q_oper):
        # TODO: deprecate and make private in favor of to_super,
        # which looks at Qobj.type to determine the right conversion function.
        """
        Takes a Choi matrix to a superoperator
        TODO: Sanitize input, Abstract-ify application of channels to states
        """
        q_oper = super_to_choi(q_oper)
        q_oper.superrep = 'super'
        return q_oper

    def choi_to_kraus(q_oper, tol=1e-9):
        """
        Takes a Choi matrix and returns a list of Kraus operators.
        TODO: Create a new class structure for quantum channels, perhaps as a
        strict sub-class of Qobj.
        """
        vals, vecs = q_oper.eigenstates(phase_fix=0)
        elements_to_keep = abs(vals) >= tol
        sqrt_vals, vecs = np.sqrt(vals[elements_to_keep]), vecs[elements_to_keep]
        shape = (np.prod(q_oper.dims[0][1]), np.prod(q_oper.dims[0][0]))
        dims = q_oper.dims[0][::-1]
        return [
            Qobj(inpt=sqrt_val * vec.data.reshape(shape, order="F"), dims=dims)
            for sqrt_val, vec in zip(sqrt_vals, vecs)
        ]

    def kraus_to_choi(kraus_ops):
        r"""
        Convert a list of Kraus operators into Choi representation of the channel.

        Essentially, kraus operators are a decomposition of a Choi matrix,
        so in this function we build Choi matrix from vector representation of Kraus operators.

        Parameters
        ----------
        kraus_ops : list[Qobj]
            The list of Kraus operators to be converted to Choi representation.

        Returns
        -------
        choi : Qobj
            A quantum object representing the same map as ``kraus_ops``, such that
            ``choi.superrep == "choi"``.
        """
        len_op = np.prod(kraus_ops[0].shape)
        # If Kraus ops have dims [M, N] in qutip notation (act on [N, N] density matrix and produce [M, M] d.m.),
        # Choi matrix Hilbert space will be [[M, N], [M, N]] because Choi Hilbert space is (output space) x (input space).
        choi_dims = [kraus_ops[0].dims] * 2
        # transform a list of Qobj matrices list[sum_ij k_ij |i><j|]
        # into an array of array vectors sum_ij k_ij |i, j>> = sum_I k_I |I>>
        kraus_vectors = np.asarray(
            [np.reshape(kraus_op.full(), len_op, "F") for kraus_op in kraus_ops]
        )
        # sum_{I} |k_I|^2 |I>><<I|
        choi_array = np.tensordot(kraus_vectors, kraus_vectors.conj(), axes=([0], [0]))
        return Qobj(choi_array, choi_dims, superrep="choi", copy=False)

    def kraus_to_super(kraus_list):
        """
        Converts a list of Kraus operators and returns a super operator.
        """
        return choi_to_super(kraus_to_choi(kraus_list))

    def _nq(dims):
        dim = np.prod(dims[0][0])
        nq = int(log2(dim))
        if 2 ** nq != dim:
            raise ValueError("{} is not an integer power of 2.".format(dim))
        return nq

    def choi_to_chi(q_oper):
        """
        Converts a Choi matrix to a Chi matrix in the Pauli basis.

        NOTE: this is only supported for qubits right now. Need to extend to
        Heisenberg-Weyl for other subsystem dimensions.
        """
        nq = _nq(q_oper.dims)
        B = _pauli_basis(nq)
        # Force the basis change to match the dimensions of
        # the input.
        B.dims = q_oper.dims
        B.superrep = 'choi'

        return Qobj(B.dag() * q_oper * B, superrep='chi')

    def chi_to_choi(q_oper):
        """
        Converts a Chi matrix to a Choi matrix.

        NOTE: this is only supported for qubits right now. Need to extend to
        Heisenberg-Weyl for other subsystem dimensions.
        """
        nq = _nq(q_oper.dims)
        B = _pauli_basis(nq)
        # Force the basis change to match the dimensions of
        # the input.
        B.dims = q_oper.dims

        # We normally should not multiply objects of different
        # superreps, so Qobj warns about that. Here, however, we're actively
        # converting between, so the superrep of B is irrelevant.
        # To suppress warnings, we pretend that B is also a chi.
        B.superrep = 'chi'

        # The Chi matrix has tr(chi) == d², so we need to divide out
        # by that to get back to the Choi form.
        return Qobj((B * q_oper * B.dag()) / q_oper.shape[0], superrep='choi')

    def _svd_u_to_kraus(U, S, d, dK, indims, outdims):
        """
        Given a partial isometry U and a vector of square-roots of singular values S
        obtained from an SVD, produces the Kraus operators represented by U.

        Returns
        -------
        Ks : list of Qobj
            Quantum objects represnting each of the Kraus operators.
        """
        # We use U * S since S is 1-index, such that this is equivalent to
        # U . diag(S), but easier to write down.
        Ks = list(map(Qobj, array(U * S).reshape((d, d, dK), order='F').transpose((2, 0, 1))))
        for K in Ks:
            K.dims = [outdims, indims]
        return Ks

    def _generalized_kraus(q_oper, thresh=1e-10):
        # TODO: document!
        # TODO: use this to generalize to_kraus to the case where U != V.
        #       This is critical for non-CP maps, as appear in (for example)
        #       diamond norm differences between two CP maps.
        if q_oper.type != "super" or q_oper.superrep != "choi":
            raise ValueError("Expected a Choi matrix, got a {} (superrep {}).".format(q_oper.type, q_oper.superrep))

        # Remember the shape of the underlying space,
        # as we'll need this to make Kraus operators later.
        dL, dR = map(int, map(sqrt, q_oper.shape))
        # Also remember the dims breakout.
        out_dims, in_dims = q_oper.dims
        out_left, out_right = out_dims
        in_left, in_right = in_dims

        # Find the SVD.
        U, S, V = svd(q_oper.full())

        # Truncate away the zero singular values, up to a threshold.
        nonzero_idxs = S > thresh
        dK = nonzero_idxs.sum()
        U = U[:, nonzero_idxs]
        S = sqrt(S[nonzero_idxs])
        # Since NumPy returns V and not V+, we need to take the dagger
        # to get back to quantum info notation for Stinespring pairs.
        V = V.conj().T[:, nonzero_idxs]

        # Next, we convert each of U and V into Kraus operators.
        # Finally, we want the Kraus index to be left-most so that we
        # can map over it when making Qobjs.
        # FIXME: does not preserve dims!
        kU = _svd_u_to_kraus(U, S, dL, dK, out_right, out_left)
        kV = _svd_u_to_kraus(V, S, dL, dK, in_right, in_left)

        return kU, kV

    def choi_to_stinespring(q_oper, thresh=1e-10):
        # TODO: document!
        kU, kV = _generalized_kraus(q_oper, thresh=thresh)

        assert(len(kU) == len(kV))
        dK = len(kU)
        dL = kU[0].shape[0]
        dR = kV[0].shape[1]
        # Also remember the dims breakout.
        out_dims, in_dims = q_oper.dims
        out_left, out_right = out_dims
        in_left, in_right = in_dims

        A = Qobj(zeros((dK * dL, dL)), dims=[out_left + [dK], out_right + [1]])
        B = Qobj(zeros((dK * dR, dR)), dims=[in_left + [dK], in_right + [1]])

        for idx_kraus, (KL, KR) in enumerate(zip(kU, kV)):
            A += tensor(KL, basis(dK, idx_kraus))
            B += tensor(KR, basis(dK, idx_kraus))

        # There is no input (right) Kraus index, so strip that off.
        del A.dims[1][-1]
        del B.dims[1][-1]

        return A, B

    # PUBLIC CONVERSION FUNCTIONS -------------------------------------------------
    # These functions handle superoperator conversions in a way that preserves the
    # correctness of Qobj.type, and in a way that automatically branches based on
    # the input Qobj.type.

    def to_choi(q_oper):
        """
        Converts a Qobj representing a quantum map to the Choi representation,
        such that the trace of the returned operator is equal to the dimension
        of the system.

        Parameters
        ----------
        q_oper : Qobj
            Superoperator to be converted to Choi representation. If
            ``q_oper`` is ``type="oper"``, then it is taken to act by conjugation,
            such that ``to_choi(A) == to_choi(sprepost(A, A.dag()))``.

        Returns
        -------
        choi : Qobj
            A quantum object representing the same map as ``q_oper``, such that
            ``choi.superrep == "choi"``.

        Raises
        ------
        TypeError: if the given quantum object is not a map, or cannot be converted
            to Choi representation.
        """
        if q_oper.type == 'super':
            if q_oper.superrep == 'choi':
                return q_oper
            if q_oper.superrep == 'super':
                return super_to_choi(q_oper)
            if q_oper.superrep == 'chi':
                return chi_to_choi(q_oper)
            else:
                raise TypeError(q_oper.superrep)
        elif q_oper.type == 'oper':
            return super_to_choi(sprepost(q_oper, q_oper.dag()))
        else:
            raise TypeError(
                "Conversion of Qobj with type = {0.type} "
                "and superrep = {0.choi} to Choi not supported.".format(q_oper)
            )

    def to_chi(q_oper):
        """
        Converts a Qobj representing a quantum map to a representation as a chi
        (process) matrix in the Pauli basis, such that the trace of the returned
        operator is equal to the dimension of the system.

        Parameters
        ----------
        q_oper : Qobj
            Superoperator to be converted to Chi representation. If
            ``q_oper`` is ``type="oper"``, then it is taken to act by conjugation,
            such that ``to_chi(A) == to_chi(sprepost(A, A.dag()))``.

        Returns
        -------
        chi : Qobj
            A quantum object representing the same map as ``q_oper``, such that
            ``chi.superrep == "chi"``.

        Raises
        ------
        TypeError: if the given quantum object is not a map, or cannot be converted
            to Chi representation.
        """
        if q_oper.type == 'super':
            # Case 1: Already done.
            if q_oper.superrep == 'chi':
                return q_oper
            # Case 2: Can directly convert.
            elif q_oper.superrep == 'choi':
                return choi_to_chi(q_oper)
            # Case 3: Need to go through Choi.
            elif q_oper.superrep == 'super':
                return to_chi(to_choi(q_oper))
            else:
                raise TypeError(q_oper.superrep)
        elif q_oper.type == 'oper':
            return to_chi(sprepost(q_oper, q_oper.dag()))
        else:
            raise TypeError(
                "Conversion of Qobj with type = {0.type} "
                "and superrep = {0.choi} to Choi not supported.".format(q_oper)
            )

    def to_super(q_oper):
        """
        Converts a Qobj representing a quantum map to the supermatrix (Liouville)
        representation.

        Parameters
        ----------
        q_oper : Qobj
            Superoperator to be converted to supermatrix representation. If
            ``q_oper`` is ``type="oper"``, then it is taken to act by conjugation,
            such that ``to_super(A) == sprepost(A, A.dag())``.

        Returns
        -------
        superop : Qobj
            A quantum object representing the same map as ``q_oper``, such that
            ``superop.superrep == "super"``.

        Raises
        ------
        TypeError
            If the given quantum object is not a map, or cannot be converted
            to supermatrix representation.
        """
        if q_oper.type == 'super':
            # Case 1: Already done.
            if q_oper.superrep == "super":
                return q_oper
            # Case 2: Can directly convert.
            elif q_oper.superrep == 'choi':
                return choi_to_super(q_oper)
            # Case 3: Need to go through Choi.
            elif q_oper.superrep == 'chi':
                return to_super(to_choi(q_oper))
            # Case 4: Something went wrong.
            else:
                raise ValueError(
                    "Unrecognized superrep '{}'.".format(q_oper.superrep))
        elif q_oper.type == 'oper':  # Assume unitary
            return sprepost(q_oper, q_oper.dag())
        else:
            raise TypeError(
                "Conversion of Qobj with type = {0.type} "
                "and superrep = {0.superrep} to supermatrix not "
                "supported.".format(q_oper)
            )

    def to_kraus(q_oper, tol=1e-9):
        """
        Converts a Qobj representing a quantum map to a list of quantum objects,
        each representing an operator in the Kraus decomposition of the given map.

        Parameters
        ----------
        q_oper : Qobj
            Superoperator to be converted to Kraus representation. If
            ``q_oper`` is ``type="oper"``, then it is taken to act by conjugation,
            such that ``to_kraus(A) == to_kraus(sprepost(A, A.dag())) == [A]``.

        tol : Float
            Optional threshold parameter for eigenvalues/Kraus ops to be discarded.
            The default is to=1e-9.

        Returns
        -------
        kraus_ops : list of Qobj
            A list of quantum objects, each representing a Kraus operator in the
            decomposition of ``q_oper``.

        Raises
        ------
        TypeError: if the given quantum object is not a map, or cannot be
            decomposed into Kraus operators.
        """
        if q_oper.type == 'super':
            if q_oper.superrep in ("super", "chi"):
                return to_kraus(to_choi(q_oper), tol)
            elif q_oper.superrep == 'choi':
                return choi_to_kraus(q_oper, tol)
        elif q_oper.type == 'oper':  # Assume unitary
            return [q_oper]
        else:
            raise TypeError(
                "Conversion of Qobj with type = {0.type} "
                "and superrep = {0.superrep} to Kraus decomposition not "
                "supported.".format(q_oper)
            )

    def to_stinespring(q_oper):
        r"""
        Converts a Qobj representing a quantum map $\Lambda$ to a pair of partial isometries
        $A$ and $B$ such that $\Lambda(X) = \Tr_2(A X B^\dagger)$ for all inputs $X$, where
        the partial trace is taken over a a new index on the output dimensions of $A$ and $B$.

        For completely positive inputs, $A$ will always equal $B$ up to precision errors.

        Parameters
        ----------
        q_oper : Qobj
            Superoperator to be converted to a Stinespring pair.

        Returns
        -------
        A, B : Qobj
            Quantum objects representing each of the Stinespring matrices for the input Qobj.
        """
        return choi_to_stinespring(to_choi(q_oper))