# Source code for qutip.random_objects

    # -*- coding: utf-8 -*-
    """
    This module is a collection of random state and operator generators.
    The sparsity of the ouput Qobj's is controlled by varing the
    `density` parameter.
    """

    __all__ = [
        'rand_herm', 'rand_unitary', 'rand_ket', 'rand_dm',
        'rand_unitary_haar', 'rand_ket_haar', 'rand_dm_ginibre',
        'rand_dm_hs', 'rand_super_bcsz', 'rand_stochastic', 'rand_super'
    ]

    import numpy as np
    import scipy.linalg as la
    import scipy.sparse as sp
    from qutip.qobj import Qobj
    from qutip.operators import create, destroy, jmat
    from qutip.states import basis
    import qutip.superop_reps as sr

    UNITS = np.array([1, 1j])

    def rand_jacobi_rotation(A, seed=None):
        """Random Jacobi rotation of a sparse matrix.

        Parameters
        ----------
        A : spmatrix
            Input sparse matrix.

        Returns
        -------
        spmatrix
            Rotated sparse matrix.
        """
        if seed is not None:
            np.random.seed(seed=seed)
        if A.shape[0] != A.shape[1]:
            raise ValueError('Input matrix must be square.')
        n = A.shape[0]
        angle = 2*np.random.random()*np.pi
        a = 1.0/np.sqrt(2)*np.exp(-1j*angle)
        b = 1.0/np.sqrt(2)*np.exp(1j*angle)
        i = int(np.floor(np.random.random()*n))
        j = i
        while i == j:
            j = int(np.floor(np.random.random()*n))
        data = np.hstack((np.array([a, -b, a, b], dtype=complex),
                          np.ones(n-2, dtype=complex)))
        diag = np.delete(np.arange(n), [i, j])
        rows = np.hstack(([i, i, j, j], diag))
        cols = np.hstack(([i, j, i, j], diag))
        R = sp.coo_matrix(
            (data, (rows, cols)), shape=(n, n), dtype=complex,
        ).tocsr()
        A = R*A*R.conj().transpose()
        return A

    def randnz(shape, norm=1 / np.sqrt(2), seed=None):
        # This function is intended for internal use.
        """
        Returns an array of standard normal complex random variates.
        The Ginibre ensemble corresponds to setting ``norm = 1`` [Mis12]_.

        Parameters
        ----------
        shape : tuple
            Shape of the returned array of random variates.
        norm : float
            Scale of the returned random variates, or 'ginibre' to draw
            from the Ginibre ensemble.
        """
        if seed is not None:
            np.random.seed(seed=seed)
        if norm == 'ginibre':
            norm = 1
        return np.sum(np.random.randn(*(shape + (2,))) * UNITS, axis=-1) * norm

    def rand_herm(N, density=0.75, dims=None, pos_def=False, seed=None):
        """Creates a random NxN sparse Hermitian quantum object.

        If 'N' is an integer, uses :math:`H=0.5*(X+X^{+})` where :math:`X` is
        a randomly generated quantum operator with a given `density`. Else uses
        complex Jacobi rotations when 'N' is given by an array.

        Parameters
        ----------
        N : int, list/ndarray
            If int, then shape of output operator. If list/ndarray then eigenvalues
            of generated operator.
        density : float
            Density between [0,1] of output Hermitian operator.
        dims : list
            Dimensions of quantum object.  Used for specifying
            tensor structure. Default is dims=[[N],[N]].
        pos_def : bool (default=False)
            Return a positive semi-definite matrix (by diagonal dominance).
        seed : int
            seed for the random number generator

        Returns
        -------
        oper : qobj
            NxN Hermitian quantum operator.

        Notes
        -----
        If given a list/ndarray as input 'N', this function returns a
        random Hermitian object with eigenvalues given in the list/ndarray.
        This is accomplished via complex Jacobi rotations.  While this method
        is ~50% faster than the corresponding (real only) Matlab code, it should
        not be repeatedly used for generating matrices larger than ~1000x1000.
        """
        if seed is not None:
            np.random.seed(seed=seed)
        if isinstance(N, (np.ndarray, list)):
            M = sp.diags(N, 0, dtype=complex, format='csr')
            N = len(N)
            if dims:
                _check_dims(dims, N, N)
            nvals = max([N**2 * density, 1])
            M = rand_jacobi_rotation(M)
            while M.nnz < 0.95 * nvals:
                M = rand_jacobi_rotation(M)
            M.sort_indices()
        elif isinstance(N, (int, np.int32, np.int64)):
            if dims:
                _check_dims(dims, N, N)
            if density < 0.5:
                M = _rand_herm_sparse(N, density, pos_def)
            else:
                M = _rand_herm_dense(N, density, pos_def)
        else:
            raise TypeError('Input N must be an integer or array_like.')
        return Qobj(M, dims=dims, isherm=True)

    def _rand_herm_sparse(N, density, pos_def):
        target = (1-(1-density)**0.5)
        num_elems = (N**2 - 0.666 * N) * target + 0.666 * N * density
        num_elems = max([num_elems, 1])
        num_elems = int(num_elems)
        data = (2 * np.random.rand(num_elems) - 1) + \
               (2 * np.random.rand(num_elems) - 1) * 1j
        row_idx, col_idx = zip(*[
            divmod(index, N)
            for index in np.random.choice(N*N, num_elems, replace=False)
        ])
        M = sp.coo_matrix((data, (row_idx, col_idx)),
                          dtype=complex, shape=(N, N))
        M = 0.5 * (M + M.conj().transpose())
        if pos_def:
            M = M.tocoo()
            M.setdiag(np.abs(M.diagonal()) + np.sqrt(2)*N)
        M = M.tocsr()
        M.sort_indices()
        return M

    def _rand_herm_dense(N, density, pos_def):
        M = (
            (2*np.random.rand(N, N) - 1)
            + 1j*(2*np.random.rand(N, N) - 1)
        )
        M = 0.5 * (M + M.conj().transpose())
        target = 1 - density**0.5
        num_remove = N * (N - 0.666) * target + 0.666 * N * (1 - density)
        num_remove = max([num_remove, 1])
        num_remove = int(num_remove)
        for index in np.random.choice(N*N, num_remove, replace=False):
            row, col = divmod(index, N)
            M[col, row] = 0
            M[row, col] = 0
        if pos_def:
            np.fill_diagonal(M, np.abs(M.diagonal()) + np.sqrt(2)*N)
        return M

    def rand_unitary(N, density=0.75, dims=None, seed=None):
        r"""Creates a random NxN sparse unitary quantum object.

        Uses :math:`\exp(-iH)` where H is a randomly generated
        Hermitian operator.

        Parameters
        ----------
        N : int
            Shape of output quantum operator.
        density : float
            Density between [0,1] of output Unitary operator.
        dims : list
            Dimensions of quantum object.  Used for specifying
            tensor structure. Default is dims=[[N],[N]].

        Returns
        -------
        oper : qobj
            NxN Unitary quantum operator.

        """
        if dims:
            _check_dims(dims, N, N)
        U = (-1.0j * rand_herm(N, density, seed=seed)).expm()
        U.data.sort_indices()
        return Qobj(U, dims=dims, shape=[N, N])

    def rand_unitary_haar(N=2, dims=None, seed=None):
        """
        Returns a Haar random unitary matrix of dimension
        ``dim``, using the algorithm of [Mez07]_.

        Parameters
        ----------
        N : int
            Dimension of the unitary to be returned.
        dims : list of lists of int, or None
            Dimensions of quantum object.  Used for specifying
            tensor structure. Default is dims=[[N],[N]].

        Returns
        -------
        U : Qobj
            Unitary of dims ``[[dim], [dim]]`` drawn from the Haar
            measure.
        """
        if dims is not None:
            _check_dims(dims, N, N)
        else:
            dims = [[N], [N]]

        # Mez01 STEP 1: Generate an N × N matrix Z of complex standard
        #               normal random variates.
        Z = randnz((N, N), seed=seed)

        # Mez01 STEP 2: Find a QR decomposition Z = Q · R.
        Q, R = la.qr(Z)

        # Mez01 STEP 3: Create a diagonal matrix Lambda by rescaling
        #               the diagonal elements of R.
        Lambda = np.diag(R).copy()
        Lambda /= np.abs(Lambda)

        # Mez01 STEP 4: Note that R' := Λ¯¹ · R has real and
        #               strictly positive elements, such that
        #               Q' = Q · Λ is Haar distributed.
        # NOTE: Λ is a diagonal matrix, represented as a vector
        #       of the diagonal entries. Thus, the matrix dot product
        #       is represented nicely by the NumPy broadcasting of
        #       the *scalar* multiplication. In particular,
        #       Q · Λ = Q_ij Λ_jk = Q_ij δ_jk λ_k = Q_ij λ_j.
        #       As NumPy arrays, Q has shape (N, N) and
        #       Lambda has shape (N, ), such that the broadcasting
        #       represents precisely Q_ij λ_j.
        U = Qobj(Q * Lambda)
        U.dims = dims
        return U

    def rand_ket(N=None, density=1, dims=None, seed=None):
        """Creates a random Nx1 sparse ket vector.

        Parameters
        ----------
        N : int
            Number of rows for output quantum vector.
            If None, N is deduced from dims.
        density : float
            Density between [0,1] of output ket state.
        dims : list
            Dimensions of quantum object.  Used for specifying
            tensor structure. Default is dims=[[N],[1]].
        seed : int
            Seed for the random number generator.

        Returns
        -------
        oper : qobj
            Nx1 ket quantum state vector.

        Raises
        -------
        ValueError
            If neither `N` or `dims` are specified.

        """
        if seed is not None:
            np.random.seed(seed=seed)
        if N is None and dims is None:
            raise ValueError('Specify either the number of rows of state vector'
                             '(N) or dimensions of quantum object (dims)')
        if N is not None and dims:
            _check_dims(dims, N, 1)
        elif dims:
            N = np.prod(dims[0])
            _check_dims(dims, N, 1)
        else:
            dims = [[N], [1]]
        X = sp.rand(N, 1, density, format='csr')
        while X.nnz == 0:
            # ensure that the ket is not all zeros.
            X = sp.rand(N, 1, density+1/N, format='csr')
        X.data = X.data - 0.5
        Y = X.copy()
        Y.data = 1.0j * (np.random.random(len(X.data)) - 0.5)
        X = X + Y
        X.sort_indices()
        X = Qobj(X)
        return Qobj(X / X.norm(), dims=dims)

    def rand_ket_haar(N=None, dims=None, seed=None):
        """
        Returns a Haar random pure state of dimension ``dim`` by
        applying a Haar random unitary to a fixed pure state.

        Parameters
        ----------
        N : int
            Dimension of the state vector to be returned.
            If None, N is deduced from dims.
        dims : list of ints, or None
            Dimensions of the resultant quantum object.
            If None, [[N],[1]] is used.

        Returns
        -------
        psi : Qobj
            A random state vector drawn from the Haar measure.

        Raises
        -------
        ValueError
            If neither `N` or `dims` are specified.
        """
        if N is None and dims is None:
            raise ValueError('Specify either the number of rows of state vector'
                             '(N) or dimensions of quantum object (dims)')
        if N and dims:
            _check_dims(dims, N, 1)
        elif dims:
            N = np.prod(dims[0])
            _check_dims(dims, N, 1)
        else:
            dims = [[N], [1]]
        psi = rand_unitary_haar(N, seed=seed) * basis(N, 0)
        psi.dims = dims
        return psi

    def rand_dm(N, density=0.75, pure=False, dims=None, seed=None):
        r"""Creates a random NxN density matrix.

        Parameters
        ----------
        N : int, ndarray, list
            If int, then shape of output operator. If list/ndarray then eigenvalues
            of generated density matrix.
        density : float
            Density between [0,1] of output density matrix.
        dims : list
            Dimensions of quantum object.  Used for specifying
            tensor structure. Default is dims=[[N],[N]].
        seed : int
            Seed for the random number generator.

        Returns
        -------
        oper : qobj
            NxN density matrix quantum operator.

        Notes
        -----
        For small density matrices., choosing a low density will result in an error
        as no diagonal elements will be generated such that :math:`Tr(\rho)=1`.

        """
        if isinstance(N, (np.ndarray, list)):
            if np.abs(np.sum(N)-1.0) > 1e-15:
                raise ValueError('Eigenvalues of a density matrix '
                                 'must sum to one.')
            H = sp.diags(N, 0, dtype=complex, format='csr')
            N = len(N)
            if dims:
                _check_dims(dims, N, N)
            nvals = N**2*density
            H = rand_jacobi_rotation(H, seed=seed)
            while H.nnz < 0.95*nvals:
                H = rand_jacobi_rotation(H)
            H.sort_indices()
        elif isinstance(N, (int, np.int32, np.int64)):
            if dims:
                _check_dims(dims, N, N)
            if pure:
                dm_density = np.sqrt(density)
                psi = rand_ket(N, dm_density, seed=seed)
                H = psi * psi.dag()
                H.data.sort_indices()
            else:
                non_zero = 0
                tries = 0
                while non_zero == 0 and tries < 10:
                    H = rand_herm(N, density, seed=seed)
                    H = H.dag() * H
                    non_zero = H.tr()
                    tries += 1
                if tries >= 10:
                    raise ValueError(
                        "Requested density is too low to generate density matrix.")
                H = H / H.tr()
                H.data.sort_indices()
        else:
            raise TypeError('Input N must be an integer or array_like.')
        return Qobj(H, dims=dims)

    def rand_dm_ginibre(N=2, rank=None, dims=None, seed=None):
        """
        Returns a Ginibre random density operator of dimension
        ``dim`` and rank ``rank`` by using the algorithm of
        [BCSZ08]_. If ``rank`` is `None`, a full-rank
        (Hilbert-Schmidt ensemble) random density operator will be
        returned.

        Parameters
        ----------
        N : int
            Dimension of the density operator to be returned.
        dims : list
            Dimensions of quantum object.  Used for specifying
            tensor structure. Default is dims=[[N],[N]].

        rank : int or None
            Rank of the sampled density operator. If None, a full-rank
            density operator is generated.

        Returns
        -------
        rho : Qobj
            An N × N density operator sampled from the Ginibre
            or Hilbert-Schmidt distribution.
        """
        if rank is None:
            rank = N
        if rank > N:
            raise ValueError("Rank cannot exceed dimension.")

        X = randnz((N, rank), norm='ginibre', seed=seed)
        rho = np.dot(X, X.T.conj())
        rho /= np.trace(rho)

        return Qobj(rho, dims=dims)

    def rand_dm_hs(N=2, dims=None, seed=None):
        """
        Returns a Hilbert-Schmidt random density operator of dimension
        ``dim`` and rank ``rank`` by using the algorithm of
        [BCSZ08]_.

        Parameters
        ----------
        N : int
            Dimension of the density operator to be returned.
        dims : list
            Dimensions of quantum object.  Used for specifying
            tensor structure. Default is dims=[[N],[N]].

        Returns
        -------
        rho : Qobj
            A dim × dim density operator sampled from the Ginibre
            or Hilbert-Schmidt distribution.

        """
        return rand_dm_ginibre(N, rank=None, dims=dims, seed=seed)

    def rand_kraus_map(N, dims=None, seed=None):
        """
        Creates a random CPTP map on an N-dimensional Hilbert space in Kraus
        form.

        Parameters
        ----------
        N : int
            Length of input/output density matrix.
        dims : list
            Dimensions of quantum object.  Used for specifying
            tensor structure. Default is dims=[[N],[N]].

        Returns
        -------
        oper_list : list of qobj
            N^2 x N x N qobj operators.

        """

        if dims:
            _check_dims(dims, N, N)

        # Random unitary (Stinespring Dilation)
        orthog_cols = rand_unitary(N ** 3, seed=seed).full()[:, :N]
        oper_list = np.reshape(orthog_cols, (N ** 2, N, N))
        return list(map(lambda x: Qobj(inpt=x, dims=dims), oper_list))

    def rand_super(N=5, dims=None, seed=None):
        """
        Returns a randomly drawn superoperator acting on operators acting on
        N dimensions.

        Parameters
        ----------
        N : int
            Square root of the dimension of the superoperator to be returned.
        dims : list
            Dimensions of quantum object.  Used for specifying
            tensor structure. Default is dims=[[[N],[N]], [[N],[N]]].
        """
        from .propagator import propagator
        if dims is not None:
            # TODO: check!
            pass
        else:
            dims = [[[N], [N]], [[N], [N]]]
        H = rand_herm(N, seed=seed)
        S = propagator(H, np.random.rand(), [
            create(N), destroy(N), jmat(float(N - 1) / 2.0, 'z')
        ])
        S.dims = dims
        return S

    def rand_super_bcsz(N=2, enforce_tp=True, rank=None, dims=None, seed=None):
        """
        Returns a random superoperator drawn from the Bruzda
        et al ensemble for CPTP maps [BCSZ08]_. Note that due to
        finite numerical precision, for ranks less than full-rank,
        zero eigenvalues may become slightly negative, such that the
        returned operator is not actually completely positive.

        Parameters
        ----------
        N : int
            Square root of the dimension of the superoperator to be returned.
        enforce_tp : bool
            If True, the trace-preserving condition of [BCSZ08]_ is enforced;
            otherwise only complete positivity is enforced.
        rank : int or None
            Rank of the sampled superoperator. If None, a full-rank
            superoperator is generated.
        dims : list
            Dimensions of quantum object.  Used for specifying
            tensor structure. Default is dims=[[[N],[N]], [[N],[N]]].

        Returns
        -------
        rho : Qobj
            A superoperator acting on vectorized dim × dim density operators,
            sampled from the BCSZ distribution.
        """
        if dims is not None:
            # TODO: check!
            pass
        else:
            dims = [[[N], [N]], [[N], [N]]]

        if rank is None:
            rank = N**2
        if rank > N**2:
            raise ValueError("Rank cannot exceed superoperator dimension.")

        # We use mainly dense matrices here for speed in low
        # dimensions. In the future, it would likely be better to switch off
        # between sparse and dense matrices as the dimension grows.

        # We start with a Ginibre uniform matrix X of the appropriate rank,
        # and use it to construct a positive semidefinite matrix X X⁺.
        X = randnz((N**2, rank), norm='ginibre', seed=seed)

        # Precompute X X⁺, as we'll need it in two different places.
        XXdag = np.dot(X, X.T.conj())

        if enforce_tp:
            # We do the partial trace over the first index by using dense reshape
            # operations, so that we can avoid bouncing to a sparse representation
            # and back.
            Y = np.einsum('ijik->jk', XXdag.reshape((N, N, N, N)))

            # Now we have the matrix 𝟙 ⊗ Y^{-1/2}, which we can find by doing
            # the square root and the inverse separately. As a possible
            # improvement, iterative methods exist to find inverse square root
            # matrices directly, as this is important in statistics.
            Z = np.kron(
                np.eye(N),
                la.sqrtm(la.inv(Y))
            )

            # Finally, we dot everything together and pack it into a Qobj,
            # marking the dimensions as that of a type=super (that is,
            # with left and right compound indices, each representing
            # left and right indices on the underlying Hilbert space).
            D = Qobj(np.dot(Z, np.dot(XXdag, Z)))
        else:
            D = N * Qobj(XXdag / np.trace(XXdag))

        D.dims = [
            # Left dims
            [[N], [N]],
            # Right dims
            [[N], [N]]
        ]

        # Since [BCSZ08] gives a row-stacking Choi matrix, but QuTiP
        # expects a column-stacking Choi matrix, we must permute the indices.
        D = D.permute([[1], [0]])

        D.dims = dims

        # Mark that we've made a Choi matrix.
        D.superrep = 'choi'

        return sr.to_super(D)

    def rand_stochastic(N, density=0.75, kind='left', dims=None, seed=None):
        """Generates a random stochastic matrix.

        Parameters
        ----------
        N : int
            Dimension of matrix.
        density : float
            Density between [0,1] of output density matrix.
        kind : str (Default = 'left')
            Generate 'left' or 'right' stochastic matrix.
        dims : list
            Dimensions of quantum object.  Used for specifying
            tensor structure. Default is dims=[[N],[N]].

        Returns
        -------
        oper : qobj
            Quantum operator form of stochastic matrix.
        """
        if seed is not None:
            np.random.seed(seed=seed)
        if dims:
            _check_dims(dims, N, N)
        num_elems = max([int(np.ceil(N*(N+1)*density)/2), N])
        data = np.random.rand(num_elems)
        # Ensure an element on every row and column
        row_idx = np.hstack([np.random.permutation(N),
                             np.random.choice(N, num_elems-N)])
        col_idx = np.hstack([np.random.permutation(N),
                             np.random.choice(N, num_elems-N)])
        M = sp.coo_matrix((data, (row_idx, col_idx)),
                          dtype=float, shape=(N, N)).tocsr()
        M = 0.5 * (M + M.conj().transpose())
        num_rows = M.indptr.shape[0]-1
        for row in range(num_rows):
            row_start = M.indptr[row]
            row_end = M.indptr[row+1]
            row_sum = np.sum(M.data[row_start:row_end])
            M.data[row_start:row_end] /= row_sum
        if kind == 'left':
            M = M.transpose()
        return Qobj(M, dims=dims, shape=(N, N))

    def _check_dims(dims, N1, N2):
        if len(dims) != 2:
            raise Exception("Qobj dimensions must be list of length 2.")
        if (not isinstance(dims[0], list)) or (not isinstance(dims[1], list)):
            raise TypeError(
                "Qobj dimension components must be lists. i.e. dims=[[N],[N]]")
        if np.prod(dims[0]) != N1 or np.prod(dims[1]) != N2:
            raise ValueError("Qobj dimensions must match matrix shape.")
        if len(dims[0]) != len(dims[1]):
            raise TypeError("Qobj dimension components must have same length.")