# Source code for qutip.core.direct_sum

    # Required for Sphinx to follow autodoc_type_aliases
    from __future__ import annotations

    from . import data as _data
    from .cy.qobjevo import QobjEvo
    from .dimensions import Dimensions, Field, SumSpace
    from .operators import qzero
    from .qobj import Qobj
    from .superoperator import operator_to_vector

    from .. import settings
    from ..typing import DimensionLike, LayerType

    from functools import partial
    from numbers import Number
    from typing import Any, Union, overload

    import numpy as np

    __all__ = ['direct_sum', 'direct_sum_sparse',
               'direct_component', 'set_direct_component']

    ARRAY_ONE = np.full((1,), 1, dtype=_data.base.idxint_dtype)

    def _is_like_ket(qobj: Any) -> bool:
        return (
            isinstance(qobj, Number)
            or (
                isinstance(qobj, (Qobj, QobjEvo))
                and qobj.type in ['scalar', 'ket']
            ))

    def _is_like_bra(qobj: Any) -> bool:
        return (
            isinstance(qobj, Number)
            or (
                isinstance(qobj, (Qobj, QobjEvo))
                and qobj.type in ['scalar', 'bra']
            ))

    def _is_like_oper(qobj: Any) -> bool:
        return (
            qobj is None
            or isinstance(qobj, Number)
            or (
                isinstance(qobj, (Qobj, QobjEvo))
                and qobj.type in ['scalar', 'oper', 'ket', 'bra']
            ))

    def _is_like_operator_ket(qobj: Any) -> bool:
        return (
            isinstance(qobj, Number)
            or (
                isinstance(qobj, (Qobj, QobjEvo))
                and qobj.type in ['scalar', 'oper', 'operator-ket']
            ))

    def _is_like_operator_bra(qobj: Any) -> bool:
        return (
            isinstance(qobj, Number)
            or (
                isinstance(qobj, (Qobj, QobjEvo))
                and qobj.type in ['scalar', 'oper', 'operator-bra']
            ))

    def _is_like_super(qobj: Any) -> bool:
        return (
            qobj is None
            or isinstance(qobj, Number)
            or (
                isinstance(qobj, (Qobj, QobjEvo))
                and qobj.type in [
                    'scalar', 'super', 'operator-ket', 'operator-bra']
            ))

    def _qobj_data(qobj: Number | Qobj, scalar_dtype: LayerType) -> _data.Data:
        if qobj is None:
            return None
        if isinstance(qobj, Number):
            return _data.identity[scalar_dtype](1, scale=qobj)
        return qobj.data

    def _qobj_dims(qobj: Number | Qobj | QobjEvo) -> Dimensions:
        return (
            qobj._dims if isinstance(qobj, (Qobj, QobjEvo))
            else Dimensions(Field(), Field())
        )

    @overload
    def direct_sum(
        qobjs: list[Number | Qobj] | list[list[Number | Qobj]],
        dtype: LayerType = None
    ) -> Qobj:
        ...

    @overload
    def direct_sum(
        qobjs: list[Number | Qobj | QobjEvo] | list[list[Number | Qobj | QobjEvo]],
        dtype: LayerType = None
    ) -> QobjEvo:
        ...

    def direct_sum(qobjs, dtype=None):
        """
        Construct the direct sum of the provided component quantum objects.
        The matrix representing the returned direct sum is a block matrix, where
        the matrices of the components are all concatenated.

        This function accepts either a 1D list of components or a 2D square matrix
        (list of lists) of components. A 1D list is interpreted as a column or row
        block vector, depending on the element types:

        * If all elements are kets or scalars, the output will be a ket.
        * If all elements are bras or scalars, the output will be a bra.
        * If all elements are operator-kets, operators, or scalars, the operators
          will be converted to operator-kets automatically, and the output will be
          an operator-ket.
        * If all elements are operator-bras, operators, or scalars, the operators
          will be converted to operator-bras automatically, and the output will be
          an operator-bra.

        Python numbers are always accepted as scalars.

        A 2D square matrix may contain any types of Qobj, but the horizontal and
        vertical dimensions must match in each block row and block column. In a 2D
        square matrix, entries may be ``None`` to indicate zero blocks. However,
        each block row and each block column must contain at least one not-``None``
        entry. For direct sums with all-zero rows or columns, consider using
        :func:`.direct_sum_sparse`.

        Parameters
        ----------
        qobjs : list
            1D list or 2D square matrix of components. The components may be Python
            numbers, :class:`.Qobj`, or :class:`.QobjEvo`.
        dtype : type or str, optional
            Storage representation for the output ``Qobj``. Any data-layer known
            to ``qutip.data.to`` is accepted.

        Returns
        -------
        direct_sum: :class:`.Qobj` or :class:`.QobjEvo`
            The assembled direct-sum object. Returns a time-dependent object
            (:class:`.QobjEvo`) if any input component is time-dependent.
        """

        if len(qobjs) == 0:
            raise ValueError("No Qobjs provided for direct sum.")

        # Step 1: check type of all provided qobj and make (row, col) -> qobj dict

        linear = isinstance(qobjs[0], (Number, Qobj, QobjEvo))
        if not linear and len(qobjs[0]) == 0:
            raise ValueError("No Qobjs provided for direct sum.")
        if not linear and not all(len(row) == len(qobjs[0]) for row in qobjs):
            raise ValueError("Matrix of Qobjs in direct sum must be square.")

        if linear and all(_is_like_ket(qobj) for qobj in qobjs):
            num_rows = len(qobjs)
            num_columns = 1
            qobj_dict = {(row, 0): qobj for row, qobj in enumerate(qobjs)}

        elif linear and all(_is_like_operator_ket(qobj) for qobj in qobjs):
            # for convenience, we call operator_to_vector on operators provided
            # in matrix form
            def _ensure_vector(qobj):
                if isinstance(qobj, (Qobj, QobjEvo)) and qobj.type == 'oper':
                    return operator_to_vector(qobj)
                return qobj

            num_rows = len(qobjs)
            num_columns = 1
            qobj_dict = {(row, 0): _ensure_vector(qobj)
                         for row, qobj in enumerate(qobjs)}

        elif linear and all(_is_like_bra(qobj) for qobj in qobjs):
            num_rows = 1
            num_columns = len(qobjs)
            qobj_dict = {(0, col): qobj for col, qobj in enumerate(qobjs)}

        elif linear and all(_is_like_operator_bra(qobj) for qobj in qobjs):
            # for convenience, we call operator_to_vector on operators provided
            # in matrix form
            def _ensure_vector(qobj):
                if isinstance(qobj, (Qobj, QobjEvo)) and qobj.type == 'oper':
                    return operator_to_vector(qobj).dag()
                return qobj

            num_rows = 1
            num_columns = len(qobjs)
            qobj_dict = {(0, col): _ensure_vector(qobj)
                         for col, qobj in enumerate(qobjs)}

        elif linear:
            raise ValueError("Invalid combination of Qobj types for direct sum.")

        else:
            allsuper = all(_is_like_super(qobj) for row in qobjs for qobj in row)
            alloper = all(_is_like_oper(qobj) for row in qobjs for qobj in row)
            if not (allsuper or alloper):
                raise ValueError(
                    "Invalid combination of Qobj types for direct sum.")

            num_rows = len(qobjs)
            num_columns = len(qobjs[0])
            qobj_dict = {(row, col): qobjs[row][col]
                         for row in range(num_rows) for col in range(num_columns)
                         if qobjs[row][col] is not None}

        # Step 2: determine dimensions of sum space

        to_spaces = [None] * num_rows
        from_spaces = [None] * num_columns

        for (row, col), qobj in qobj_dict.items():
            to_space = _qobj_dims(qobj)[0]
            if to_spaces[row] is None:
                to_spaces[row] = to_space
            elif to_spaces[row] != to_space:
                raise ValueError(
                    f"Direct sum: inconsistent dimensions in row {row}. Expected"
                    f" {to_spaces[row].as_list()}, got {to_space.as_list()}.")

            from_space = _qobj_dims(qobj)[1]
            if from_spaces[col] is None:
                from_spaces[col] = from_space
            elif from_spaces[col] != from_space:
                raise ValueError(
                    f"Direct sum: inconsistent dimensions in column {col}."
                    f" Expected {from_spaces[col].as_list()}, got"
                    f" {from_space.as_list()}.")

        for row, space in enumerate(to_spaces):
            if space is None:
                raise ValueError(f"Direct sum: empty row {row}.")
        for col, space in enumerate(from_spaces):
            if space is None:
                raise ValueError(f"Direct sum: empty column {col}.")

        sum_dimension = Dimensions(SumSpace(from_spaces), SumSpace(to_spaces))

        # Step 3: create direct sum

        return _do_direct_sum(qobj_dict, sum_dimension, dtype)

    @overload
    def direct_sum_sparse(
        qobjs: dict[tuple[int, int], Number | Qobj],
        sum_dimensions: DimensionLike,
        dtype: LayerType = None
    ) -> Qobj:
        ...

    @overload
    def direct_sum_sparse(
        qobjs: dict[tuple[int, int], Number | Qobj | QobjEvo],
        sum_dimensions: DimensionLike,
        dtype: LayerType = None
    ) -> QobjEvo:
        ...

    def direct_sum_sparse(qobjs, sum_dimensions, dtype=None):
        r"""
        Construct the direct sum of the provided component quantum objects.
        This is a variant of :func:`.direct_sum` suitable for large, sparse block
        matrices. The caller must provide the dimensions of the direct sum, for
        example in list form as in :code:`[([2], [3]), ([2], [3])]` for an
        operator on :math:`\mathbb C^2 \oplus \mathbb C^3`.

        In more detail: the ``sum_dimensions`` parameter is a list of two entries.
        The two entries of this list respectively correspond to the output and
        input Hilbert spaces of the quantum object. Each entry is the "list
        representation" of the dimensions of the corresponding Hilbert space, like
        it is used throughout QuTiP. For use in this function, either one of the
        two spaces or both of them must have the structure of a direct sum, which
        is represented as a tuple in the "list representation".

        List representation syntax overview:

        * A normal, unmarked Hilbert space of dimension :math:`n` is represented
          by the list :code:`[n]`.
        * A product Hilbert space, where the constituents have the dimensions
          :math:`n_1` up to :math:`n_k`, is represented by the list
          :code:`[n1, n2, ..., nk]`.
        * The direct sum of Hilbert spaces, which have the list representations
          :code:`l1` up to :code:`[lk]`, is represented by the tuple
          :code:`[l1, l2, ..., lk]`. Sum dimensions may be nested.
        * For example, the direct sum of a 2D space and a 3D space has the list
          representation :code:`([2], [3])`.
        * As another example, the direct sum of a 2D space, and the product of two
          3D spaces, is :code:`([2], [3, 3])`.
        * Superoperators act on operator spaces; the list representation of an
          operator space is the list :code:`[lOut, lIn]`, where :code:`lIn` and
          :code:`lOut` are list representations of the input and output spaces
          of the operators. Direct sum components may be either regular Hilbert
          spaces or operator spaces, but not mixed within the same quantum object.

        Parameters
        ----------
        qobjs : dict
            Mapping that assigns a component to the given (row, col) location in
            the block matrix. Missing locations are filled with zero blocks. The
            components may be Python numbers, :class:`.Qobj`, or :class:`.QobjEvo`.
        sum_dimensions : list of tuples or lists
            Dimensions of the resulting Qobj, like in the creation of a Qobj.
            See detailed explanation above.
        dtype : type or str, optional
            Storage representation for the output ``Qobj``. Any data-layer known
            to ``qutip.data.to`` is accepted.

        Returns
        -------
        direct_sum: :class:`.Qobj` or :class:`.QobjEvo`
            The assembled direct-sum object. Returns a time-dependent object
            (:class:`.QobjEvo`) if any input component is time-dependent.
        """

        sum_dimensions = Dimensions(sum_dimensions)
        to_spaces, from_spaces = _spaces_from_dims(sum_dimensions)

        for (row, col), qobj in qobjs.items():
            _check_bounds(row, 0, len(to_spaces))
            _check_bounds(col, 0, len(from_spaces))
            if (
                _qobj_dims(qobj)[0] != to_spaces[row] or
                _qobj_dims(qobj)[1] != from_spaces[col]
            ):
                raise ValueError("Direct sum: dimension mismatch for component at"
                                 f" ({row}, {col}).")

        return _do_direct_sum(dict(sorted(qobjs.items())), sum_dimensions, dtype)

    def _do_direct_sum(
        qobjs: dict[tuple[int, int], Number | Qobj | QobjEvo],
        sum_dimensions: Dimensions,
        dtype: LayerType = None
    ):
        """Assumes `qobjs` is sorted and performs no dimensions checks"""

        if settings.core["default_dtype_scope"] == "full":
            dtype = _data._parse_default_dtype(dtype, "sparse")
        else:
            dtype = _data._parse_default_dtype(dtype or "core", "sparse")

        nonevo_count = sum(
            not isinstance(qobj, QobjEvo) for qobj in qobjs.values()
        )
        block_rows = np.empty((nonevo_count,), dtype=_data.base.idxint_dtype)
        block_cols = np.empty((nonevo_count,), dtype=_data.base.idxint_dtype)
        blocks = np.empty((nonevo_count,), dtype=_data.Data)

        i = 0
        qobjevos = []
        for (row, col), qobj in qobjs.items():
            if not isinstance(qobj, QobjEvo):
                block_rows[i] = row
                block_cols[i] = col
                blocks[i] = _qobj_data(qobj, dtype)
                i += 1
            else:
                qobjevos.append((row, col, qobj))

        if isinstance(sum_dimensions[0], SumSpace):
            block_heights = np.array(
                sum_dimensions[0]._space_dims, dtype=_data.base.idxint_dtype)
        else:
            block_heights = np.full(
                (1,), sum_dimensions[0].size, dtype=_data.base.idxint_dtype)
        if isinstance(sum_dimensions[1], SumSpace):
            block_widths = np.array(
                sum_dimensions[1]._space_dims, dtype=_data.base.idxint_dtype)
        else:
            block_widths = np.full(
                (1,), sum_dimensions[1].size, dtype=_data.base.idxint_dtype)

        result = Qobj(
            _data.block_build(
                block_rows, block_cols, blocks,
                block_heights, block_widths, dtype=dtype
            ), dims=sum_dimensions, copy=False
        )

        for row, col, qobjevo in qobjevos:
            result += qobjevo.linear_map(
                partial(_blow_up_qobj, sum_dimensions=sum_dimensions,
                        row=row, col=col, dtype=dtype),
                _skip_check=True
            )
        if qobjevos:
            result.compress(_skip_coeff=True)
        return result

    @overload
    def direct_component(sum_qobj: Qobj, *index: int | slice) -> Qobj:
        ...

    @overload
    def direct_component(sum_qobj: QobjEvo, *index: int | slice) -> Qobj | QobjEvo:
        ...

    def direct_component(sum_qobj, *index):
        """
        Extract the component (block) at the given index from a direct sum
        ``Qobj``.

        Parameters
        ----------
        sum_qobj : :class:`.Qobj` or :class:`.QobjEvo`
            A direct sum object from which to extract a component.
        index : list of int or slice
            One or two indices identifying the component to extract. If both the
            row and column spaces are sums, two indices (``row``, ``col``) are
            required. If only one side is a sum, a single index is accepted and
            interpreted as the index into the summed side.

        Returns
        -------
        component: :class:`.Qobj` or :class:`.QobjEvo`
            The extracted component.

        Notes
        -----
        Each index may be either an int indicating a single component or a slice
        indicating multiple components. For convenience, slices can be created
        using `np.s_[start:stop]`. Custom steps are not supported.
        """

        if settings.core["default_dtype_scope"] == "full":
            dtype = settings.core["default_dtype"] or sum_qobj.dtype
        else:
            dtype = sum_qobj.dtype

        if isinstance(sum_qobj, QobjEvo):
            result = sum_qobj.linear_map(lambda x: direct_component(x, *index),
                                         _skip_check=True)
            result.compress(_skip_coeff=True)
            return result(0) if result.isconstant else result

        to_index, from_index = _check_component_index(sum_qobj._dims, index)
        component_dims, row_start, row_stop, col_start, col_stop =\
            _component_info(sum_qobj._dims, to_index, from_index)

        out_data = _data.block_extract(
            sum_qobj.data, row_start, row_stop, col_start, col_stop, dtype=dtype
        )
        return Qobj(out_data, dims=component_dims, copy=False)

    @overload
    def set_direct_component(
        sum_qobj: Qobj,
        component: Qobj,
        *index: int | slice,
        dtype: LayerType = None
    ) -> Qobj:
        ...

    @overload
    def set_direct_component(
        sum_qobj: Qobj | QobjEvo,
        component: Qobj | QobjEvo,
        *index: int | slice,
        dtype: LayerType = None
    ) -> QobjEvo:
        ...

    def set_direct_component(sum_qobj, component, *index):
        """
        Set (replace) a component in a direct sum ``Qobj``. The function returns a
        new object where the component at the given index is replaced with
        ``component``.

        Parameters
        ----------
        sum_qobj : :class:`.Qobj` or :class:`.QobjEvo`
            The direct sum object whose component will be replaced.
        component : :class:`.Qobj` or :class:`.QobjEvo`
            The new component to insert. ``None`` sets the component to zero.
        index : list of int or slice
            One or two indices identifying the component to extract. If both the
            row and column spaces are sums, two indices (``row``, ``col``) are
            required. If only one side is a sum, a single index is accepted and
            interpreted as the index into the summed side.

        Returns
        -------
        updated: :class:`.Qobj` or :class:`.QobjEvo`
            The resulting direct sum object with the component set.

        Notes
        -----
        Each index may be either an int indicating a single component or a slice
        indicating multiple components. For convenience, slices can be created
        using `np.s_[start:stop]`. Custom steps are not supported.
        """

        if settings.core["default_dtype_scope"] == "full":
            dtype = settings.core["default_dtype"] or sum_qobj.dtype
        else:
            dtype = sum_qobj.dtype

        to_index, from_index = _check_component_index(sum_qobj._dims, index)
        component_dims, row_start, _, col_start, _ =\
            _component_info(sum_qobj._dims, to_index, from_index)

        if component is not None and _qobj_dims(component) != component_dims:
            raise ValueError("Canot set component of direct sum: dimension"
                             f" mismatch. Expected: {component_dims.as_list()},"
                             f" got: {_qobj_dims(component).as_list()}.")

        if not (isinstance(sum_qobj, QobjEvo) or isinstance(component, QobjEvo)):
            component_data = (
                _qobj_data(component, dtype)
                or _data.zeros[dtype](*component_dims.shape)
            )
            out_data = _data.block_overwrite(sum_qobj.data, component_data,
                                             row_start, col_start, dtype=dtype)
            return Qobj(out_data, dims=sum_qobj._dims, copy=False)

        # If QobjEvo is involved, the result is the sum of two parts:
        # - the sum_qobj with the component zeroed
        # - the new component embedded into a large zero matrix
        # Note that all "set_component" below act on simple Qobj

        if isinstance(sum_qobj, QobjEvo):
            zeroed = sum_qobj.linear_map(
                lambda x: set_direct_component(x, None, *index),
                _skip_check=True)
            zeroed.compress(_skip_coeff=True)
            zeroed = zeroed(0) if zeroed.isconstant else zeroed
        else:
            zeroed = set_direct_component(sum_qobj, None, *index)

        blow_up_func = partial(_blow_up_qobj, sum_dimensions=sum_qobj._dims,
                               row=to_index, col=from_index, dtype=dtype)
        if isinstance(component, QobjEvo):
            blow_up = component.linear_map(blow_up_func, _skip_check=True)
        else:
            blow_up = blow_up_func(component)

        return zeroed + blow_up

    def _blow_up_qobj(x: Qobj, *, sum_dimensions, row, col, dtype):
        # given small Qobj x, makes big Qobj with given dimensions
        # resulting big Qobj is zero except for the block x starting at row, col

        if x is None:
            return qzero(sum_dimensions[0], sum_dimensions[1])
        if type(row) is slice:
            row = row.start
        if type(col) is slice:
            col = col.start

        if isinstance(sum_dimensions[0], SumSpace):
            data_row = sum_dimensions[0]._space_cumdim(row)
        else:
            data_row = 0

        if isinstance(sum_dimensions[1], SumSpace):
            data_col = sum_dimensions[1]._space_cumdim(col)
        else:
            data_col = 0

        xh, xw = x.shape
        sh, sw = sum_dimensions.shape

        return Qobj(
            _data.block_build(
                ARRAY_ONE,
                ARRAY_ONE,
                np.full((1,), x.data, dtype=_data.Data),
                np.array([data_row, xh, sh - xh - data_row],
                         dtype=_data.base.idxint_dtype),
                np.array([data_col, xw, sw - xw - data_col],
                         dtype=_data.base.idxint_dtype),
                dtype=dtype
            ),
            dims=sum_dimensions, copy=False
        )

    def _check_bounds(given, min, max):
        if type(given) is slice:
            if given.step and given.step != 1:
                raise IndexError("Slicing with custom step is not supported.")
            if given.start >= given.stop:
                raise IndexError("Empty slice is not supported.")
            if given.start < min or given.stop > max:
                raise IndexError(f"Index [{given.start}:{given.stop}] out of"
                                 f" bounds ({min}, {max-1}) for component of"
                                 " direct sum")
        elif not (min <= given < max):
            raise IndexError(f"Index ({given}) out of bounds ({min}, {max-1})"
                             " for component of direct sum.")

    def _spaces_from_dims(sum_dimensions):
        if isinstance(sum_dimensions[0], SumSpace):
            to_spaces = sum_dimensions[0].spaces
        else:
            to_spaces = [sum_dimensions[0]]
        if isinstance(sum_dimensions[1], SumSpace):
            from_spaces = sum_dimensions[1].spaces
        else:
            from_spaces = [sum_dimensions[1]]
        return to_spaces, from_spaces

    def _check_component_index(sum_dimensions, index):
        """
        Check that an appropriate number of indices is provided for the given
        dimensions. 2 indices is always okay, 1 index is also okay if only one of
        the spaces is a sum. Returns tuple (to_index, from_index) of 2 indices.
        """
        to_spaces, from_spaces = _spaces_from_dims(sum_dimensions)
        is_to_sum = len(to_spaces) > 1
        is_from_sum = len(from_spaces) > 1

        if is_to_sum and is_from_sum:
            if len(index) != 2:
                raise IndexError("Invalid number of indices provided for component"
                                 " of direct sum (two indices required).")
        else:
            if len(index) == 1 and is_to_sum:
                index = (index[0], 0)
            elif len(index) == 1:
                index = (0, index[0])
            if len(index) != 2:
                raise IndexError("Invalid number of indices provided for component"
                                 " of direct sum (one or two indices required).")

        return index

    def _component_info(sum_dimensions, to_index, from_index):
        """
        For the component (block) indicated by the indices, return
        * the Dimensions of the block
        * the index of the first row of the block in the full data array
        * the index of the row below the block in the full data array
        * the index of the first column of the block in the full data array
        * the index of the column after the block in the full data array
        """
        if isinstance(sum_dimensions[0], SumSpace):
            _check_bounds(to_index, 0, len(sum_dimensions[0].spaces))
            component_to = sum_dimensions[0].spaces[to_index]
            if type(to_index) is slice:
                data_row_start = sum_dimensions[0]._space_cumdim(to_index.start)
                data_row_stop = sum_dimensions[0]._space_cumdim(to_index.stop)
            else:
                data_row_start = sum_dimensions[0]._space_cumdim(to_index)
                data_row_stop = sum_dimensions[0]._space_cumdim(to_index + 1)
        else:
            _check_bounds(to_index, 0, 1)
            component_to = sum_dimensions[0]
            data_row_start = 0
            data_row_stop = sum_dimensions[0].size
        if type(to_index) is slice:
            component_to = SumSpace(component_to)

        if isinstance(sum_dimensions[1], SumSpace):
            _check_bounds(from_index, 0, len(sum_dimensions[1].spaces))
            component_from = sum_dimensions[1].spaces[from_index]
            if type(from_index) is slice:
                data_col_start = sum_dimensions[1]._space_cumdim(from_index.start)
                data_col_stop = sum_dimensions[1]._space_cumdim(from_index.stop)
            else:
                data_col_start = sum_dimensions[1]._space_cumdim(from_index)
                data_col_stop = sum_dimensions[1]._space_cumdim(from_index + 1)
        else:
            _check_bounds(from_index, 0, 1)
            component_from = sum_dimensions[1]
            data_col_start = 0
            data_col_stop = sum_dimensions[1].size
        if type(from_index) is slice:
            component_from = SumSpace(component_from)

        return (Dimensions(component_from, component_to),
                data_row_start, data_row_stop,
                data_col_start, data_col_stop)