# !/usr/bin/env python3
# Copyright (c) 2023 QuAIR team. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

r"""
The source file of the Circuit class.
"""

import math
import shutil
import warnings
from typing import Callable, Dict, Iterable, List, Optional, Tuple, Union

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import torch

from .ansatz import (ComplexBlockLayer, ComplexEntangledLayer, Layer,
                     LinearEntangledLayer, OperatorList, RealBlockLayer,
                     RealEntangledLayer, TrotterLayer, Universal2, Universal3)
from .core import Hamiltonian, intrinsic, latex, utils
from .core.intrinsic import _alias
from .operator import (CCX, CNOT, CP, CRX, CRY, CRZ, CSWAP, CU, CY, CZ, MS, RX,
                       RXX, RY, RYY, RZ, RZZ, SWAP, U3, AmplitudeDamping,
                       BitFlip, BitPhaseFlip, ChoiRepr, Collapse,
                       ControlOracle, ControlParamOracle, Depolarizing, Gate,
                       GeneralizedAmplitudeDamping, GeneralizedDepolarizing, H,
                       KrausRepr, OneWayLOCC, Oracle, P, ParamOracle,
                       PauliChannel, Permutation, PhaseDamping, PhaseFlip,
                       ResetChannel, S, Sdg, StinespringRepr, T, Tdg,
                       ThermalRelaxation, UniversalQudits, X, Y, Z)
from .operator.gate import _circuit_plot

__all__ = ['Circuit']


class Circuit(OperatorList):
    r"""Class for quantum circuit.

    Args:
        num_systems: number of systems in the circuit. Defaults to None. Alias of ``num_qubits``.
        system_dim: dimension of systems of this circuit. Can be a list of system dimensions 
            or an int representing the dimension of all systems. Defaults to be qubit case.
        physical_idx: physical indices of systems. Defaults to be the same as the logical indices.
    
    Note:
        when the number of system is unknown and system_dim is an int, the circuit is a dynamic quantum circuit.

    Examples:
        .. code-block:: python

            qc = Circuit(1)  # A quantum circuit with 1 qubit
            qc.h()
            print(f'The latex code of this circuit is:\n{qc.to_latex()}')

        ::

            The latex code of this circuit is:
            \lstick{} & \gate[1]{H}
    """
    @_alias({'num_systems': 'num_qubits'})
    def __init__(self, num_systems: Optional[int] = None, 
                 system_dim: Optional[Union[List[int], int]] = 2,
                 physical_idx: Optional[List[int]] = None) -> None:
        super().__init__(num_systems, system_dim, physical_idx)

        # alias
        self.toffoli = self.ccx
        self.cx = self.cnot
        self.collapse = self.measure

    def _get_drawer(self, style: str, decimal: int) -> Tuple[latex.OperatorListDrawer, Dict[int, str]]:
        r"""Return the drawer that draws the circuit in LaTeX format, and the beginning string of the circuit.

        Examples:
            .. code-block:: python

                drawer, begin_code = qc._Circuit__get_drawer('standard', 2)

            ::

                drawer: (an instance of latex.OperatorListDrawer)
                begin_code: {0: r'\lstick{}', 1: r'\lstick{}', ...}
        """
        physical_idx, self.system_idx = self.system_idx, list(range(self.num_systems))
        
        drawer = latex.OperatorListDrawer(style, decimal)
        
        for op in self.children():
            if isinstance(op, Layer):
                drawer = drawer + op._get_drawer(style, decimal)
            else:
                drawer.append(op.info)
        
        begin_code = {idx: r'\lstick{}' for idx in self.system_idx}
        self.system_idx = physical_idx
        return drawer, begin_code
        
    def to_latex(self, style: str = 'standard', decimal: int = 2) -> str:
        r"""The LaTeX representation of the circuit, written in Quantikz format.
        
        Args:
            style: the style of the plot, can be 'standard', 'compact' or 'detailed'. Defaults to ``standard``.
            decimal: number of decimal places to display. Defaults to 2.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.h()
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{H} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{H} & {} & {}
        """
        physical_idx, self.system_idx = self.system_idx, list(range(self.num_systems))
        drawer, begin_code = self._get_drawer(style, decimal)
        drawer.fill_all()
        self.system_idx = physical_idx
        return latex.code_to_str(drawer.code, begin_code)
    
    def plot(self, style: str = 'standard', decimal: int = 2,
             dpi: int = 300, print_code: bool = False, 
             show_plot: bool = True, include_empty: bool = False,
             latex: bool = True, **kwargs) -> Optional[matplotlib.figure.Figure]:
        r'''Display the circuit using Quantikz if ``latex`` is True, otherwise using matplotlib.

        Args:
            style: the style of the plot, can be 'standard', 'compact' or 'detailed'. Defaults to ``standard``.
            decimal: number of decimal places to display. Defaults to 2.
            dpi: dots per inches of plot image. Defaults to 300.
            print_code: whether print the LaTeX code of the circuit, default to ``False``.
            show_plot: whether show the plotted circuit, default to ``True``.
            include_empty: whether include empty lines, default to ``False``.
            latex: whether use Quantikz, a LaTeX package, to plot circuits , default to ``True``.
            kwargs: additional parameters for matplotlib plot.
        
        Returns:
            None, or a ``matplotlib.figure.Figure`` instance depending on ``latex`` and ``output``.

        Notes:
            If ``latex`` is True, the circuit will be displayed in LaTeX format;
            if ``latex`` is False, the circuit will be displayed in matplotlib format,
            in which case we have three additional parameters:
            - output_plot: whether output the plot instance, default to ``False``.
            - save_path: the save path of image. Defaults to None.
            - scale: scale coefficient of figure. Default to 1.0.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.h()
                qc.plot()

            ::

                (Displays the plotted circuit using Quantikz or matplotlib)
                
        .. warning::
        
            Starting from QuAIRKit 2.5.0, `Circuit.plot` now defaults to using QuantiKz, 
            a LaTeX package powered by TikZ, for rendering quantum circuit diagrams 
            commonly used in academic publications. Support for plotting circuits with 
            Matplotlib will be deprecated in the near future.

            To fully utilize this feature, please ensure that a TeX distribution, such 
            as [TeX Live](https://www.tug.org/texlive), is installed on your system. 
            This will enhance your experience with QuAIRKit and quantum computing 
            visualization.
        '''
        if latex and (shutil.which("pdflatex") is None):
            warnings.warn(
                "pdflatex is mot detected on your system. Will skip the plot.", UserWarning)
            
            if print_code:
                print(self.to_latex(style, decimal))
            return

        if latex:
            drawer, begin_code = self._get_drawer(style, decimal)   
            drawer._fill_empty(list(range(self.num_systems if include_empty else max(drawer._code.keys()))))
            drawer.add_end()
            _fig = drawer.plot(dpi, print_code, begin_code)
            
            if show_plot:
                intrinsic._display_png(_fig)
        else:
            warnings.warn(
                "Starting from QuAIRKit 0.4.0, `Circuit.plot` now defaults to using Quantikz, a LaTeX package "
                "powered by TikZ, for rendering quantum circuit diagrams commonly used in academic publications. "
                "Support for plotting circuits with Matplotlib will be deprecated in the near future. \n"
                "To fully utilize this feature, please ensure that a TeX distribution, such as TeX Live "
                "(https://www.tug.org/texlive), is installed on your system. This will enhance your experience "
                "with QuAIRKit and quantum computing visualization.", FutureWarning)
            save_path, scale = kwargs.get('save_path', None), kwargs.get('scale', 1.0)
            output_plot = kwargs.get('output_plot', False)
            
            physical_idx, self.system_idx = self.system_idx, list(range(self.num_systems))
            _fig = _circuit_plot(self, dpi=dpi, scale=scale)
            self.system_idx = physical_idx
            
            if save_path:
                plt.savefig(save_path, dpi=dpi)
            if show_plot:
                plt.show()
            if output_plot:
                return _fig


    # ---------------------- below are common operators ---------------------- #
    def h(self, qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add single-qubit Hadamard gates.

        The matrix form of such a gate is:

        .. math::
            H = \frac{1}{\sqrt{2}}
                \begin{bmatrix}
                    1 &  1 \\
                    1 & -1
                \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'full'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.h()
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{H} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{H} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(H(qubits_idx))

    def s(self, qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add single-qubit S gates.

        The matrix form of such a gate is:

        .. math::
            S =
                \begin{bmatrix}
                    1 & 0 \\
                    0 & i
                \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'full'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.s()
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{S} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{S} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(S(qubits_idx))

    def sdg(self, qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add single-qubit S dagger (S inverse) gates.

        The matrix form of such a gate is:

        .. math::
            S^\dagger =
                \begin{bmatrix}
                    1 &  0 \\
                    0 & -i
                \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to ``'full'``.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.sdg()
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{S^\dagger} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{S^\dagger} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(Sdg(qubits_idx))

    def t(self, qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add single-qubit T gates.

        The matrix form of such a gate is:

        .. math::
            T = \begin{bmatrix}
                    1 & 0 \\
                    0 & e^{i\pi/4}
                \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'full'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.t()
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{T} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{T} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(T(qubits_idx))

    def tdg(self, qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add single-qubit T dagger (T inverse) gates.

        The matrix form of such a gate is:

        .. math::
            T^\dagger = \begin{bmatrix}
                            1 & 0 \\
                            0 & e^{-i\pi/4}
                         \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to ``'full'``.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.tdg()
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{T^\dagger} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{T^\dagger} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(Tdg(qubits_idx))

    def x(self, qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add single-qubit X gates.

        The matrix form of such a gate is:

        .. math::
           X = \begin{bmatrix}
                    0 & 1 \\
                    1 & 0
               \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'full'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.x()
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{X} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{X} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(X(qubits_idx))

    def y(self, qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add single-qubit Y gates.

        The matrix form of such a gate is:

        .. math::
            Y = \begin{bmatrix}
                    0 & -i \\
                    i &  0
                \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'full'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.y()
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{Y} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{Y} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(Y(qubits_idx))

    def z(self, qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add single-qubit Z gates.

        The matrix form of such a gate is:

        .. math::
            Z = \begin{bmatrix}
                    1 &  0 \\
                    0 & -1
                \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'full'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.z()
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{Z} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{Z} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(Z(qubits_idx))

    def p(self, qubits_idx: Union[Iterable[int], int, str] = 'full',
          param: Union[torch.Tensor, float] = None, param_sharing: bool = False) -> None:
        r"""Add single-qubit P gates.

        The matrix form of such a gate is:

        .. math::
            P(\theta) = \begin{bmatrix}
                            1 & 0 \\
                            0 & e^{i\theta}
                        \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'full'.
            param: Parameters of the gates. Defaults to None.
            param_sharing: Whether gates in the same layer share a parameter. Defaults to False.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.p([0], torch.pi/2)
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{P(1.57)} & \meter[2]{} & {} \\
                \lstick{} & {} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(P(qubits_idx, param, param_sharing))

    def rx(self, qubits_idx: Union[Iterable[int], int, str] = 'full',
           param: Union[torch.Tensor, float] = None, param_sharing: bool = False) -> None:
        r"""Add single-qubit rotation gates about the x-axis.

        The matrix form of such a gate is:

        .. math::
            R_X(\theta) = \begin{bmatrix}
                                \cos\frac{\theta}{2} & -i\sin\frac{\theta}{2} \\
                                -i\sin\frac{\theta}{2} & \cos\frac{\theta}{2}
                           \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'full'.
            param: Parameters of the gates. Defaults to None.
            param_sharing: Whether gates in the same layer share a parameter. Defaults to False.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.rx([0], torch.pi/2)
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{R_{x}(1.57)} & \meter[2]{} & {} \\
                \lstick{} & {} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(RX(qubits_idx, param, param_sharing))

    def ry(self, qubits_idx: Union[Iterable[int], int, str] = 'full',
           param: Union[torch.Tensor, float] = None, param_sharing: bool = False) -> None:
        r"""Add single-qubit rotation gates about the y-axis.

        The matrix form of such a gate is:

        .. math::
            R_Y(\theta) = \begin{bmatrix}
                                \cos\frac{\theta}{2} & -\sin\frac{\theta}{2} \\
                                \sin\frac{\theta}{2} & \cos\frac{\theta}{2}
                           \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'full'.
            param: Parameters of the gates. Defaults to None.
            param_sharing: Whether gates in the same layer share a parameter. Defaults to False.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.ry([0], torch.pi/2)
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{R_{y}(1.57)} & \meter[2]{} & {} \\
                \lstick{} & {} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(RY(qubits_idx, param, param_sharing))

    def rz(self, qubits_idx: Union[Iterable[int], int, str] = 'full',
           param: Union[torch.Tensor, float] = None, param_sharing: bool = False) -> None:
        r"""Add single-qubit rotation gates about the z-axis.

        The matrix form of such a gate is:

        .. math::
            R_Z(\theta) = \begin{bmatrix}
                                e^{-i\theta/2} & 0 \\
                                0 & e^{i\theta/2}
                           \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'full'.
            param: Parameters of the gates. Defaults to None.
            param_sharing: Whether gates in the same layer share a parameter. Defaults to False.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.rz([0], torch.pi/2)
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{R_{z}(1.57)} & \meter[2]{} & {} \\
                \lstick{} & {} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(RZ(qubits_idx, param, param_sharing))

    def u3(self, qubits_idx: Union[Iterable[int], int, str] = 'full',
           param: Union[torch.Tensor, Iterable[float]] = None, param_sharing: bool = False) -> None:
        r"""Add single-qubit rotation gates.

        The matrix form of such a gate is:

        .. math::
            U_3(\theta, \phi, \lambda) =
                \begin{bmatrix}
                    \cos\frac\theta2 & -e^{i\lambda}\sin\frac\theta2 \\
                    e^{i\phi}\sin\frac\theta2 & e^{i(\phi+\lambda)}\cos\frac\theta2
                \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'full'.
            param: Parameters of the gates. Defaults to None.
            param_sharing: Whether gates in the same layer share a parameter. Defaults to False.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.u3([0], torch.pi/2)
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{U(1.57, 1.57, 1.57)} & \meter[2]{} & {} \\
                \lstick{} & {} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(U3(qubits_idx, param, param_sharing))

    def cnot(self, qubits_idx: Union[Iterable[int], str] = 'cycle') -> None:
        r"""Add CNOT gates.

        For a 2-qubit circuit, when `qubits_idx` is `[0, 1]`, the matrix form is:

        .. math::
            \mathit{CNOT} = |0\rangle \langle 0|\otimes I + |1\rangle \langle 1|\otimes X

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'cycle'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.cnot([0, 1])
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \ctrl[]{1} & \meter[2]{} & {} \\
                \lstick{} & \targ{} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 2)
        self.append(CNOT(qubits_idx))

    def cy(self, qubits_idx: Union[Iterable[int], str] = 'cycle') -> None:
        r"""Add controlled Y gates.

        For a 2-qubit circuit, when `qubits_idx` is `[0, 1]`, the matrix form is:

        .. math::
            \mathit{CY} = |0\rangle \langle 0|\otimes I + |1\rangle \langle 1|\otimes Y

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'cycle'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.cy([0, 1])
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \ctrl[]{1} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{Y} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 2)
        self.append(CY(qubits_idx))

    def cz(self, qubits_idx: Union[Iterable[int], str] = 'linear') -> None:
        r"""Add controlled Z gates.

        For a 2-qubit circuit, when `qubits_idx` is `[0, 1]`, the matrix form is:

        .. math::
            \mathit{CZ} = |0\rangle \langle 0|\otimes I + |1\rangle \langle 1|\otimes Z

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'linear'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.cz([0, 1])
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \ctrl[]{1} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{Z} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 2)
        self.append(CZ(qubits_idx))

    def swap(self, qubits_idx: Union[Iterable[int], str] = 'linear') -> None:
        r"""Add SWAP gates.

        The matrix form is:

        .. math::
            \mathit{SWAP} =
            \begin{bmatrix}
                1 & 0 & 0 & 0 \\
                0 & 0 & 1 & 0 \\
                0 & 1 & 0 & 0 \\
                0 & 0 & 0 & 1
            \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'linear'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.swap([0, 1])
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[2,style={draw=none}]{\permute{2,1}} & \meter[2]{} & {} \\
                \lstick{} & {} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 2)
        self.append(SWAP(qubits_idx))

    def cp(self, qubits_idx: Union[Iterable[int], str] = 'cycle',
           param: Union[torch.Tensor, float] = None, param_sharing: bool = False) -> None:
        r"""Add controlled P gates.

        For a 2-qubit circuit, when `qubits_idx` is `[0, 1]`, the matrix form is:

        .. math::
            \mathit{CP}(\theta) =
            \begin{bmatrix}
                1 & 0 & 0 & 0 \\
                0 & 1 & 0 & 0 \\
                0 & 0 & 1 & 0 \\
                0 & 0 & 0 & e^{i\theta}
            \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'cycle'.
            param: Parameters of the gates. Defaults to None.
            param_sharing: Whether gates in the same layer share a parameter. Defaults to False.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.cp([0, 1], torch.pi / 2)
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \ctrl[]{1} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{P(1.57)} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 2)
        self.append(CP(qubits_idx, param, param_sharing))

    def crx(self, qubits_idx: Union[Iterable[int], str] = 'cycle',
            param: Union[torch.Tensor, float] = None, param_sharing: bool = False) -> None:
        r"""Add controlled rotation gates about the x-axis.

        For a 2-qubit circuit, when `qubits_idx` is `[0, 1]`, the matrix form is:

        .. math::
            \mathit{CR_X} =
            |0\rangle \langle 0|\otimes I + |1\rangle \langle 1|\otimes R_X

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'cycle'.
            param: Parameters of the gates. Defaults to None.
            param_sharing: Whether gates in the same layer share a parameter. Defaults to False.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.crx([0, 1], torch.pi / 2)
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \ctrl[]{1} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{R_{x}(1.57)} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 2)
        self.append(CRX(qubits_idx, param, param_sharing))

    def cry(self, qubits_idx: Union[Iterable[int], str] = 'cycle',
            param: Union[torch.Tensor, float] = None, param_sharing: bool = False) -> None:
        r"""Add controlled rotation gates about the y-axis.

        For a 2-qubit circuit, when `qubits_idx` is `[0, 1]`, the matrix form is:

        .. math::
            \mathit{CR_Y} =
            |0\rangle \langle 0|\otimes I + |1\rangle \langle 1|\otimes R_Y

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'cycle'.
            param: Parameters of the gates. Defaults to None.
            param_sharing: Whether gates in the same layer share a parameter. Defaults to False.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.cry([0, 1], torch.pi / 2)
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \ctrl[]{1} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{R_{y}(1.57)} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 2)
        self.append(CRY(qubits_idx, param, param_sharing))

    def crz(self, qubits_idx: Union[Iterable[int], str] = 'cycle',
            param: Union[torch.Tensor, float] = None, param_sharing: bool = False) -> None:
        r"""Add controlled rotation gates about the z-axis.

        For a 2-qubit circuit, when `qubits_idx` is `[0, 1]`, the matrix form is:

        .. math::
            \mathit{CR_Z} =
            |0\rangle \langle 0|\otimes I + |1\rangle \langle 1|\otimes R_Z

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'cycle'.
            param: Parameters of the gates. Defaults to None.
            param_sharing: Whether gates in the same layer share a parameter. Defaults to False.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.crz([0, 1], torch.pi / 2)
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \ctrl[]{1} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{R_{z}(1.57)} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 2)
        self.append(CRZ(qubits_idx, param, param_sharing))

    def cu(self, qubits_idx: Union[Iterable[int], str] = 'cycle',
           param: Union[torch.Tensor, float] = None, param_sharing: bool = False) -> None:
        r"""Add controlled single-qubit rotation gates.

        For a 2-qubit circuit, when `qubits_idx` is `[0, 1]`, the matrix form is given by a controlled-U gate.

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'cycle'.
            param: Parameters of the gate. Defaults to None.
            param_sharing: Whether gates in the same layer share a parameter. Defaults to False.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.cu([0, 1], torch.pi / 2)
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \ctrl[]{1} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{U(1.57, ...)} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 2)
        self.append(CU(qubits_idx, param, param_sharing))

    def rxx(self, qubits_idx: Union[Iterable[int], str] = 'linear',
            param: Union[torch.Tensor, float] = None, param_sharing: bool = False) -> None:
        r"""Add RXX gates.

        The matrix form is:

        .. math::
            R_{XX}(\theta) =
            \begin{bmatrix}
                \cos\frac{\theta}{2} & 0 & 0 & -i\sin\frac{\theta}{2} \\
                0 & \cos\frac{\theta}{2} & -i\sin\frac{\theta}{2} & 0 \\
                0 & -i\sin\frac{\theta}{2} & \cos\frac{\theta}{2} & 0 \\
                -i\sin\frac{\theta}{2} & 0 & 0 & \cos\frac{\theta}{2}
            \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'linear'.
            param: Parameters of the gate. Defaults to None.
            param_sharing: Whether gates in the same layer share a parameter. Defaults to False.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.rxx([0, 1], torch.pi / 2)
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[2]{R_{xx}(1.57)} & \meter[2]{} & {} \\
                \lstick{} & {} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 2)
        self.append(RXX(qubits_idx, param, param_sharing))

    def ryy(self, qubits_idx: Union[Iterable[int], str] = 'linear',
            param: Union[torch.Tensor, float] = None, param_sharing: bool = False) -> None:
        r"""Add RYY gates.

        The matrix form is:

        .. math::
            R_{YY}(\theta) =
            \begin{bmatrix}
                \cos\frac{\theta}{2} & 0 & 0 & i\sin\frac{\theta}{2} \\
                0 & \cos\frac{\theta}{2} & -i\sin\frac{\theta}{2} & 0 \\
                0 & -i\sin\frac{\theta}{2} & \cos\frac{\theta}{2} & 0 \\
                i\sin\frac{\theta}{2} & 0 & 0 & \cos\frac{\theta}{2}
            \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'linear'.
            param: Parameters of the gate. Defaults to None.
            param_sharing: Whether gates in the same layer share a parameter. Defaults to False.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.ryy([0, 1], torch.pi / 2)
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[2]{R_{yy}(1.57)} & \meter[2]{} & {} \\
                \lstick{} & {} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 2)
        self.append(RYY(qubits_idx, param, param_sharing))

    def rzz(self, qubits_idx: Union[Iterable[int], str] = 'linear',
            param: Union[torch.Tensor, float] = None, param_sharing: bool = False) -> None:
        r"""Add RZZ gates.

        The matrix form is:

        .. math::
            R_{ZZ}(\theta) =
            \begin{bmatrix}
                e^{-i\theta/2} & 0 & 0 & 0 \\
                0 & e^{i\theta/2} & 0 & 0 \\
                0 & 0 & e^{i\theta/2} & 0 \\
                0 & 0 & 0 & e^{-i\theta/2}
            \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'linear'.
            param: Parameters of the gate. Defaults to None.
            param_sharing: Whether gates in the same layer share a parameter. Defaults to False.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.rzz([0, 1], torch.pi / 2)
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[2]{R_{zz}(1.57)} & \meter[2]{} & {} \\
                \lstick{} & {} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 2)
        self.append(RZZ(qubits_idx, param, param_sharing))

    def ms(self, qubits_idx: Union[Iterable[int], str] = 'cycle') -> None:
        r"""Add Mølmer-Sørensen (MS) gates.

        The matrix form is:

        .. math::
            \mathit{MS} = R_{XX}(-\pi/2) = \frac{1}{\sqrt{2}}
            \begin{bmatrix}
                1 & 0 & 0 & i \\
                0 & 1 & i & 0 \\
                0 & i & 1 & 0 \\
                i & 0 & 0 & 1
            \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'cycle'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.ms([0, 1])
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[2]{\text{MS}} & \meter[2]{} & {} \\
                \lstick{} & {} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 2)
        self.append(MS(qubits_idx))

    def cswap(self, qubits_idx: Union[Iterable[int], str] = 'cycle') -> None:
        r"""Add CSWAP (Fredkin) gates.

        The matrix form is:

        .. math::
            \mathit{CSWAP} =
            \begin{bmatrix}
                1 & 0 & \cdots & 0 \\
                0 & 1 &         & 0 \\
                \vdots &  & \ddots  & \vdots \\
                0 & 0 & \cdots & 1
            \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'cycle'.

        Examples:
            .. code-block:: python

                qc = Circuit(3)
                qc.cswap([0, 1, 2])
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \ctrl[]{1} & \meter[3]{} & {} \\
                \lstick{} & \gate[2,style={draw=gray, dashed}]{\permute{2,1}} & {} & {} \\
                \lstick{} & {} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 2)
        self.append(CSWAP(qubits_idx))

    def ccx(self, qubits_idx: Union[Iterable[int], str] = 'cycle') -> None:
        r"""Add CCX (Toffoli) gates.

        The matrix form is:

        .. math::
            \mathit{CCX} =
            \begin{bmatrix}
                1 & 0 & \cdots & 0 \\
                0 & 1 &         & 0 \\
                \vdots &  & \ddots  & \vdots \\
                0 & 0 & \cdots & 1
            \end{bmatrix}

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied. Defaults to 'cycle'.

        Examples:
            .. code-block:: python

                qc = Circuit(3)
                qc.ccx([0, 1, 2])
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \ctrl[]{2} & \meter[3]{} & {} \\
                \lstick{} & \control{} & {} & {} \\
                \lstick{} & \targ{} & {} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 3)
        self.append(CCX(qubits_idx))

    def universal_two_qubits(self, qubits_idx: Union[List[int], str] = None,
                             param: Union[torch.Tensor, float] = None) -> None:
        r"""Add universal two-qubit gates. One such gate requires 15 parameters.

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied.
            param: Parameters of the gates. Defaults to None.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.universal_two_qubits([0, 1])
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                (The output latex code shows a universal 2-qubit layer.)
        """
        qubits_idx = self.register_idx(qubits_idx, None)
        self.append(Universal2(qubits_idx, param))

    def universal_three_qubits(self, qubits_idx: Optional[List[int]] = None,
                               param: Union[torch.Tensor, float] = None) -> None:
        r"""Add universal three-qubit gates. One such gate requires 81 parameters.

        Args:
            qubits_idx: Indices of the qubits on which the gates are applied.
            param: Parameters of the gates. Defaults to None.

        Examples:
            .. code-block:: python

                qc = Circuit(3)
                qc.universal_three_qubits([0, 1, 2])
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                (The output latex code shows a universal 3-qubit layer.)
        """
        qubits_idx = self.register_idx(qubits_idx, None)
        self.append(Universal3(qubits_idx, param))
    
    @_alias({"system_idx": "qubits_idx"})
    def universal_qudits(self, system_idx: List[int],
                         param: Union[torch.Tensor, float] = None, param_sharing: bool = False) -> None:
        r"""Add universal qudit gates. One such gate requires :math:`d^2 - 1` parameters,
        where :math:`d` is the gate dimension.

        Args:
            system_idx: Indices of the systems on which the gates are applied.
            param: Parameters of the gates. Defaults to None.
            param_sharing: Whether gates in the same layer share a parameter. Defaults to False.

        Examples:
            .. code-block:: python

                qc = Circuit(1, 3)
                qc.universal_qudits([0])
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{\operatorname{UNI}_{3}} & \meter[1]{} & {}
        """
        system_idx = self.register_idx(system_idx, 1 if isinstance(system_idx, int) else len(system_idx))
        self.append(UniversalQudits(system_idx, [self.system_dim[idx] for idx in system_idx],
                                    param, param_sharing))
    
    @_alias({"system_idx": "qubits_idx"})
    def permute(self, perm: List[int], system_idx: List[int]) -> None:
        r"""Add a permutation gate.

        Args:
            perm: A list representing the permutation of subsystems.
            system_idx: Indices of the systems on which the gates are applied.

        Examples:
            .. code-block:: python

                qc = Circuit(3)
                qc.permute([1, 0, 2], [0, 1, 2])
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[3,style={draw=none}]{\permute{2,1,3}} & \meter[3]{} & {} \\
                \lstick{} & {} & {} & {} \\
                \lstick{} & {} & {} & {}
        """
        system_idx = self.register_idx(system_idx, len(system_idx))
        self.append(Permutation(perm, system_idx, [self.system_dim[idx] for idx in system_idx]))
        
    @_alias({"system_idx": "qubits_idx"})
    def control_permute(self, perm: List[int], system_idx: List[int], control_idx: int = -1) -> None:
        r"""Add a controlled permutation gate.

        Args:
            perm: A list representing the permutation of subsystems.
            system_idx: Indices of the systems on which the gates are applied.
            control_idx: the index that controls the oracle. Defaults to -1, meaning the highest index.

        Examples:
            .. code-block:: python

                qc = Circuit(4)
                qc.control_permute([1, 0, 2], [3, 0, 1, 2], 1)
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[3,style={draw=gray, dashed}]{\permute{2,1,3}} & \meter[4]{} & {} \\
                \lstick{} & {} & {} & {} \\
                \lstick{} & {} & {} & {} \\
                \lstick{} & \ctrl[]{-1} & {} & {}
        """
        if isinstance(system_idx[0], int):
            system_idx[0] = [system_idx[0]]
        _system_idx = system_idx[0] + system_idx[1:]
        _system_idx = self.register_idx(_system_idx, len(_system_idx))
        acted_system_dim = [self.system_dim[idx] for idx in _system_idx]

        ctrl_dim = math.prod(acted_system_dim[:len(system_idx[0])])
        if control_idx == -1:
            control_idx = ctrl_dim - 1
        else:
            assert control_idx < ctrl_dim, \
                f"Control index out of range: expect < {ctrl_dim}, got {control_idx}."
        
        permutation = utils.matrix._permutation(perm, acted_system_dim[len(system_idx[0]):])
        
        gate_info = {
            "name": "ctrl-permute",
            "api": "control_permute",
            "permute": perm,
            'plot_width': 0.2,
        }
        self.append(ControlOracle(
            permutation, system_idx, control_idx, acted_system_dim, gate_info))
    
    @_alias({"system_idx": "qubits_idx"})
    def oracle(self, oracle: torch.Tensor, system_idx: Union[List[Union[List[int], int]], int], 
               control_idx: Optional[int] = None, gate_name: Optional[str] = 'oracle',
               latex_name: Optional[str] = None) -> None:
        r"""Add an oracle gate.

        Args:
            oracle: Unitary oracle.
            system_idx: Indices of the systems on which the gate is applied.
            control_idx: the index that controls the oracle. Defaults to None.
            gate_name: name of the oracle.
            latex_name: LaTeX name of the gate. Defaults to gate_name.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.oracle(oracle=eye(2), system_idx=[0], latex_name=r'$Identity$')
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{$Identity$}
        """
        gate_info = {
            'name': gate_name,
            'tex': r'\text{' + gate_name + r'}' if latex_name is None else latex_name
        }
        
        if control_idx is None:
            system_idx = [system_idx] if isinstance(system_idx, int) else system_idx
            system_idx = self.register_idx(system_idx, len(system_idx))
            
            acted_system_dim = [self.system_dim[idx] for idx in system_idx]
            self.append(Oracle(oracle, system_idx, acted_system_dim, gate_info))
        else:
            if isinstance(system_idx[0], int):
                system_idx[0] = [system_idx[0]]
            _system_idx = system_idx[0] + system_idx[1:]
            _system_idx = self.register_idx(_system_idx, len(_system_idx))
            acted_system_dim = [self.system_dim[idx] for idx in _system_idx]

            ctrl_dim = math.prod(acted_system_dim[:len(system_idx[0])])
            if control_idx == -1:
                control_idx = ctrl_dim - 1
            else:
                assert control_idx < ctrl_dim, \
                    f"Control index out of range: expect < {ctrl_dim}, got {control_idx}."
            self.append(ControlOracle(
                oracle, system_idx, control_idx, acted_system_dim, gate_info))

    @_alias({"system_idx": "qubits_idx"})
    def control_oracle(self, oracle: torch.Tensor, system_idx: List[Union[List[int], int]], control_idx: int = -1,
                        gate_name: Optional[str] = 'oracle', latex_name: Optional[str] = None) -> None:
        r"""Add a controlled oracle gate.

        Args:
            oracle: Unitary oracle.
            system_idx: Indices of the systems on which the gate is applied.
            control_idx: the index that controls the oracle. Defaults to -1.
            gate_name: name of the oracle.
            latex_name: LaTeX name of the gate.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.control_oracle(oracle=eye(2), system_idx=[0, 1], control_idx=0, latex_name=r'$Identity$')
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \octrl[]{1} & \meter[2]{} & {} \\
                \lstick{} & \gate[1]{$Identity$} & {} & {}
        """
        self.oracle(oracle, system_idx, control_idx, gate_name, latex_name)
        
    @_alias({"system_idx": "qubits_idx"})
    def param_oracle(self, generator: Callable[[torch.Tensor], torch.Tensor], num_acted_param: int,
                     system_idx: Union[List[int], int], control_idx: Optional[int] = None,
                     param: Union[torch.Tensor, float] = None, gate_name: Optional[str] = 'param-oracle',
                     latex_name: Optional[str] = None, support_batch: bool = False) -> None:
        r"""Add a parameterized oracle gate.

        Args:
            generator: Function to generate the oracle.
            num_acted_param: Number of parameters required.
            system_idx: Indices of the systems on which the gate acts.
            control_idx: The index that controls the oracle. Defaults to None.
            param: Input parameters for the gate. Defaults to None.
            gate_name: Name of the oracle.
            latex_name: LaTeX name of the gate.
            support_batch: Whether generator supports batched input.

        Examples:
            .. code-block:: python

                def rotation_generator(params: torch.Tensor) -> torch.Tensor:
                    theta = params[..., 0]
                    cos_theta = torch.cos(theta / 2).unsqueeze(-1)
                    sin_theta = torch.sin(theta / 2).unsqueeze(-1)
                    matrix = torch.cat([
                        torch.cat([cos_theta, -1j * sin_theta], dim=-1),
                        torch.cat([-1j * sin_theta, cos_theta], dim=-1)
                    ], dim=-2)
                    return matrix

                qc = Circuit(2)
                qc.param_oracle(
                    generator=rotation_generator,
                    num_acted_param=1,
                    system_idx=[0, 1],
                    control_idx=0,
                    param=None,
                    gate_name="ControlledRotation",
                    support_batch=True
                )
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \octrl[]{1} \\
                \lstick{} & \gate[1]{ControlledRotation(0.36)}
        """
        gate_info = {
            'name': gate_name,
            'tex': r'\text{' + gate_name + r'}' if latex_name is None else latex_name}
        
        if control_idx is None:
            system_idx = [system_idx] if isinstance(system_idx, int) else system_idx
            system_idx = self.register_idx(system_idx, len(system_idx))
            
            acted_system_dim = [self.system_dim[idx] for idx in system_idx]
            self.append(ParamOracle(
                generator, system_idx, param, num_acted_param, acted_system_dim, gate_info, support_batch))
        else:
            if isinstance(system_idx[0], int):
                system_idx[0] = [system_idx[0]]
            _system_idx = system_idx[0] + system_idx[1:]
            _system_idx = self.register_idx(_system_idx, len(_system_idx))
            acted_system_dim = [self.system_dim[idx] for idx in _system_idx]
            
            ctrl_dim = math.prod(acted_system_dim[:len(system_idx[0])])
            assert control_idx < ctrl_dim, f"Control index out of range: expect < {ctrl_dim}, got {control_idx}."
            self.append(ControlParamOracle(
                generator, system_idx, control_idx, param, num_acted_param, acted_system_dim, gate_info, support_batch))

    @_alias({"system_idx": "qubits_idx", "post_selection": "desired_result"})
    def measure(self, system_idx: Union[Iterable[int], int, str] = None, post_selection: Union[int, str] = None,
                if_print: bool = False, measure_basis: Optional[torch.Tensor] = None) -> None:
        r"""Perform a measurement on the specified systems.

        Args:
            system_idx: Systems to measure. Defaults to all.
            post_selection: The post-selection result. Defaults to None.
            if_print: Whether to print collapse info. Defaults to False.
            measure_basis: Measurement basis.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.measure()
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \meter[2]{} & {} \\
                \lstick{} & {} & {}
        """
        system_idx = self.register_idx(system_idx, None)
        acted_system_dim = [self.system_dim[idx] for idx in system_idx]
        self.append(Collapse(system_idx, acted_system_dim, post_selection, if_print, measure_basis))
    
    @_alias({"system_idx": "qubits_idx"})
    def locc(self, local_unitary: torch.Tensor, system_idx: List[Union[List[int], int]],
             label: str = 'M', latex_name: str = 'O') -> None:
        r"""Add a one-way LOCC protocol comprised of unitary operations.

        Args:
            local_unitary: The local unitary operation.
            system_idx: Systems on which the protocol is applied. The first element indicates the measure system.
            label: Label for measurement. Defaults to 'M'.
            latex_name: LaTeX name for the applied operator. Defaults to 'O'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.locc(local_unitary=x(), system_idx=[0, 1])
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \meter{} & \push{M} \wireoverride{c} & \ctrl[vertical wire=c]{1}\wireoverride{c}  \\
                \lstick{} & {} & {} & \gate[1]{O^{(M)}}
        
        """
        if isinstance(system_idx[0], int):
            system_idx[0] = [system_idx[0]]
        measure_idx, apply_idx = system_idx[0], system_idx[1:]
        _system_idx = measure_idx + apply_idx
        _system_idx = self.register_idx(_system_idx, len(np.unique(_system_idx)))
        system_dim = [self.system_dim[idx] for idx in _system_idx]
        measure_dim, apply_dim = system_dim[:len(measure_idx)], system_dim[len(measure_idx):]
        
        gate = Oracle(local_unitary, apply_idx, apply_dim)
        self.append(OneWayLOCC(gate, measure_idx, measure_dim, label=label, latex_name=latex_name))
        
    @_alias({"system_idx": "qubits_idx"})
    def param_locc(self, generator: Callable[[torch.Tensor], torch.Tensor], num_acted_param: int, 
                   system_idx: List[Union[List[int], int]], param: Union[torch.Tensor, float] = None, 
                   label: str = 'M', latex_name: str = 'U', support_batch: bool = False) -> None:
        r"""Add a one-way LOCC protocol comprised of unitary operations, where the applied unitary is parameterized.

        Args:
            generator: Function to generate the oracle.
            num_acted_param: Number of parameters required.
            system_idx: Systems on which the protocol is applied. The first element indicates the measure system.
            param: Input parameters for the gate. Defaults to None.
            label: Label for measurement. Defaults to 'M'.
            latex_name: LaTeX name for the applied operator. Defaults to 'U'.
            support_batch: Whether generator supports batched input.

        """
        if isinstance(system_idx[0], int):
            system_idx[0] = [system_idx[0]]
        measure_idx, apply_idx = system_idx[0], system_idx[1:]
        _system_idx = measure_idx + apply_idx
        _system_idx = self.register_idx(_system_idx, len(np.unique(_system_idx)))
        system_dim = [self.system_dim[idx] for idx in _system_idx]
        measure_dim, apply_dim = system_dim[:len(measure_idx)], system_dim[len(measure_idx):]
        
        if param is None:
            float_dtype = intrinsic._get_float_dtype(self.dtype)
            expect_shape = intrinsic._format_param_shape([apply_idx], num_acted_param, param_sharing=False, batch_size=np.prod(measure_dim))
            param = torch.nn.Parameter(torch.rand(expect_shape, dtype=float_dtype) * 2 * np.pi)
        param_gate = ParamOracle(
            generator, apply_idx, param, num_acted_param, apply_dim, support_batch=support_batch)
        self.append(OneWayLOCC(param_gate, measure_idx, measure_dim, label=label, latex_name=latex_name))

    def linear_entangled_layer(self, qubits_idx: Optional[List[int]] = None, depth: int = 1,
                               param: Union[torch.Tensor, float] = None) -> None:
        r"""Add linear entangled layers consisting of Ry gates, Rz gates, and CNOT gates.

        Args:
            qubits_idx: Systems to apply the layer on. Defaults to all.
            depth: Number of layers. Defaults to 1.
            param: Parameters for the layer. Defaults to self-generated.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.linear_entangled_layer([0, 1], depth=4)
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                (The output latex code shows a linear entangled layer.)
        """
        qubits_idx = self.register_idx(qubits_idx, None)
        self.append(LinearEntangledLayer(qubits_idx, depth, param))

    def real_entangled_layer(self, qubits_idx: Optional[List[int]] = None, depth: int = 1,
                             param: Union[torch.Tensor, float] = None) -> None:
        r"""Add strongly entangled layers consisting of Ry gates and CNOT gates.

        Args:
            qubits_idx: Systems to apply the layer on. Defaults to all.
            depth: Number of layers. Defaults to 1.
            param: Layer parameters. Defaults to self-generated.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.real_entangled_layer([0, 1], depth=4)
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                (The output latex code shows a real entangled layer.)
        """
        qubits_idx = self.register_idx(qubits_idx, None)
        self.append(RealEntangledLayer(qubits_idx, depth, param))

    def complex_entangled_layer(self, qubits_idx: Optional[List[int]] = None, depth: int = 1,
                                param: Union[torch.Tensor, float] = None) -> None:
        r"""Add strongly entangled layers consisting of single-qubit rotation gates and CNOT gates.

        Args:
            qubits_idx: Systems to apply the layer on. Defaults to all.
            depth: Number of layers. Defaults to 1.
            param: Layer parameters. Defaults to self-generated.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.complex_entangled_layer([0, 1], depth=4)
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                (The output latex code shows a complex entangled layer.)
        """
        qubits_idx = self.register_idx(qubits_idx, None)
        self.append(ComplexEntangledLayer(qubits_idx, depth, param))

    def real_block_layer(self, qubits_idx: Optional[List[int]] = None, depth: int = 1,
                         param: Union[torch.Tensor, float] = None) -> None:
        r"""Add weakly entangled layers consisting of Ry gates and CNOT gates.

        Args:
            qubits_idx: Systems to apply the layer on. Defaults to all.
            depth: Number of layers. Defaults to 1.
            param: Layer parameters. Defaults to self-generated.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.real_block_layer([0, 1], depth=4)
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                (The output latex code shows a real block layer.)
        """
        qubits_idx = self.register_idx(qubits_idx, None)
        self.append(RealBlockLayer(qubits_idx, depth, param))

    def complex_block_layer(self, qubits_idx: Optional[List[int]] = None, depth: int = 1,
                            param: Union[torch.Tensor, float] = None) -> None:
        r"""Add weakly entangled layers consisting of single-qubit rotation gates and CNOT gates.

        Args:
            qubits_idx: Systems to apply the layer on. Defaults to all.
            depth: Number of layers. Defaults to 1.
            param: Layer parameters. Defaults to self-generated.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.complex_block_layer([0, 1], depth=4)
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                (The output latex code shows a complex block layer.)
        """
        qubits_idx = self.register_idx(qubits_idx, None)
        self.append(ComplexBlockLayer(qubits_idx, depth, param))

    def trotter(
        self, hamiltonian: Hamiltonian, time: float, qubits_idx: Optional[List[int]] = None, 
        num_steps: int = 1, order: int = 1, name: str = 'H'
    ) -> None:
        r"""Add Trotter decompositions of a Hamiltonian evolution operator.

        Args:
            hamiltonian: Hamiltonian of the system whose time evolution is to be simulated.
            time: Total evolution time.
            qubits_idx: Indices of the qubits on which the layer is applied. Defaults to ``None``.
            num_steps: Number of trotter blocks. Defaults to 1.
            order: Order of the Trotter-Suzuki decomposition. Defaults to 1.
            name: Name of the Hamiltonian. Defaults to 'H'.

        """
        qubits_idx = self.register_idx(qubits_idx, None)
        tau = time / num_steps
        self.append(TrotterLayer(hamiltonian, qubits_idx, tau, num_steps, order, name))

    def bit_flip(
            self, prob: Union[torch.Tensor, float],
            qubits_idx: Union[Iterable[int], int, str] = 'full'
    ) -> None:
        r"""Add bit flip channels.

        Args:
            prob: Probability of a bit flip.
            qubits_idx: Systems to apply the channel on. Defaults to 'full'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.bit_flip(prob=0.5, qubits_idx=[0, 1])
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{\mathcal{E}_{p = 0.5}^{\textrm{\tiny{(BF)}}}} \\
                \lstick{} & \gate[1]{\mathcal{E}_{p = 0.5}^{\textrm{\tiny{(BF)}}}}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(BitFlip(prob, qubits_idx))

    def phase_flip(self, prob: Union[torch.Tensor, float],
                   qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add phase flip channels.

        Args:
            prob: Probability of a phase flip.
            qubits_idx: Systems to apply the channel on. Defaults to 'full'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.phase_flip(prob=0.5, qubits_idx=[0, 1])
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{\mathcal{E}_{p = 0.5}^{\textrm{\tiny{(PF)}}}} \\
                \lstick{} & \gate[1]{\mathcal{E}_{p = 0.5}^{\textrm{\tiny{(PF)}}}}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(PhaseFlip(prob, qubits_idx))

    def bit_phase_flip(self, prob: Union[torch.Tensor, float],
                       qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add bit phase flip channels.

        Args:
            prob: Probability of a bit phase flip.
            qubits_idx: Systems to apply the channel on. Defaults to 'full'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.bit_phase_flip(prob=0.5, qubits_idx=[0, 1])
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{\mathcal{E}_{p = 0.5}^{\textrm{\tiny{(BPF)}}}} \\
                \lstick{} & \gate[1]{\mathcal{E}_{p = 0.5}^{\textrm{\tiny{(BPF)}}}}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(BitPhaseFlip(prob, qubits_idx))

    def amplitude_damping(self, gamma: Union[torch.Tensor, float],
                          qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add amplitude damping channels.

        Args:
            gamma: Damping probability.
            qubits_idx: Systems to apply the damping on. Defaults to 'full'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.amplitude_damping(gamma=0.5, qubits_idx=[0, 1])
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{\mathcal{E}_{\gamma = 0.5}^{\textrm{\tiny{(AD)}}}} \\
                \lstick{} & \gate[1]{\mathcal{E}_{\gamma = 0.5}^{\textrm{\tiny{(AD)}}}}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(AmplitudeDamping(gamma, qubits_idx))

    def generalized_amplitude_damping(self, gamma: Union[torch.Tensor, float],
                                      prob: Union[torch.Tensor, float],
                                      qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add generalized amplitude damping channels.

        Args:
            gamma: Damping probability.
            prob: Excitation probability.
            qubits_idx: Systems to apply the channel on. Defaults to 'full'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.generalized_amplitude_damping(gamma=0.5, prob=0.5, qubits_idx=[0, 1])
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{\mathcal{E}_{\gamma = 0.5, p = 0.5}^{\textrm{\tiny{(GAD)}}}} \\
                \lstick{} & \gate[1]{\mathcal{E}_{\gamma = 0.5, p = 0.5}^{\textrm{\tiny{(GAD)}}}}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(GeneralizedAmplitudeDamping(gamma, prob, qubits_idx))

    def phase_damping(self, gamma: Union[torch.Tensor, float],
                      qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add phase damping channels.

        Args:
            gamma: Phase damping parameter.
            qubits_idx: Systems to apply the channel on. Defaults to 'full'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.phase_damping(gamma=0.5, qubits_idx=[0, 1])
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{\mathcal{E}_{\gamma = 0.5}^{\textrm{\tiny{(PD)}}}} \\
                \lstick{} & \gate[1]{\mathcal{E}_{\gamma = 0.5}^{\textrm{\tiny{(PD)}}}}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(PhaseDamping(gamma, qubits_idx))

    def depolarizing(self, prob: Union[torch.Tensor, float],
                     qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add depolarizing channels.

        Args:
            prob: Depolarizing probability.
            qubits_idx: Systems to apply the channel on. Defaults to 'full'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.depolarizing(prob=0.5, qubits_idx=[0, 1])
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{\mathcal{D}_{p = 0.5}} \\
                \lstick{} & \gate[1]{\mathcal{D}_{p = 0.5}}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(Depolarizing(prob, qubits_idx))

    def generalized_depolarizing(self, prob: Union[torch.Tensor, float],
                                  qubits_idx: Union[Iterable[int], int, str]) -> None:
        r"""Add a general depolarizing channel.

        Args:
            prob: Probabilities for the Pauli basis.
            qubits_idx: Systems to apply the channel on.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.generalized_depolarizing(prob=0.5, qubits_idx=[0, 1])
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[2]{\mathcal{D}_{p = 0.5}} \\
                \lstick{} & {}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(GeneralizedDepolarizing(prob, qubits_idx))

    def pauli_channel(self, prob: Union[torch.Tensor, float],
                      qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add Pauli channels.

        Args:
            prob: Probabilities for the Pauli X, Y, and Z operators.
            qubits_idx: Systems to apply the channel on. Defaults to 'full'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.pauli_channel(prob=[0.1, 0.3, 0.5])
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{\mathcal{N}} \\
                \lstick{} & \gate[1]{\mathcal{N}}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(PauliChannel(prob, qubits_idx))

    def reset_channel(self, prob: Union[torch.Tensor, float],
                      qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add reset channels.

        Args:
            prob: Probabilities for resetting to the basis states.
            qubits_idx: Systems to apply the channel on. Defaults to 'full'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.reset_channel(prob=[0.5, 0.4], qubits_idx=[0, 1])
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{\mathcal{N}} \\
                \lstick{} & \gate[1]{\mathcal{N}}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(ResetChannel(prob, qubits_idx))

    def thermal_relaxation(self, const_t: Union[torch.Tensor, Iterable[float]], exec_time: Union[torch.Tensor, float],
                           qubits_idx: Union[Iterable[int], int, str] = 'full') -> None:
        r"""Add thermal relaxation channels.

        Args:
            const_t: The T1 and T2 relaxation times.
            exec_time: Gate execution time.
            qubits_idx: Systems to apply the channel on. Defaults to 'full'.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.thermal_relaxation(const_t=[600, 300], exec_time=500, qubits_idx=[0, 1])
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{\mathcal{N}} \\
                \lstick{} & \gate[1]{\mathcal{N}}
        """
        qubits_idx = self.register_idx(qubits_idx, 1)
        self.append(ThermalRelaxation(const_t, exec_time, qubits_idx))

    def choi_channel(self, choi_repr: Iterable[torch.Tensor],
                     system_idx: Union[Iterable[Iterable[int]], Iterable[int], int]) -> None:
        r"""Add custom channels in the Choi representation.

        Args:
            choi_repr: Choi representation.
            system_idx: Systems to apply the channel on.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                X = torch.tensor(x(), dtype=torch.complex64)
                choi = X.kron(X) / 2
                qc.choi_channel(choi_repr=choi, system_idx=[0])
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{\mathcal{N}}
        """
        system_idx = [system_idx] if isinstance(system_idx, int) else system_idx
        system_idx = self.register_idx(system_idx, len(system_idx))
        acted_system_dim = [self.system_dim[idx] for idx in system_idx]
        self.append(ChoiRepr(choi_repr, system_idx, acted_system_dim))

    def kraus_channel(self, kraus_oper: Iterable[torch.Tensor],
                      system_idx: Union[Iterable[Iterable[int]], Iterable[int], int]) -> None:
        r"""Add custom channels in the Kraus representation.

        Args:
            kraus_oper: Kraus operators.
            system_idx: Systems to apply the channel on.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.kraus_channel(kraus_oper=eye(2), system_idx=[0])
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{\mathcal{N}}
        """
        system_idx = [system_idx] if isinstance(system_idx, int) else system_idx
        system_idx = self.register_idx(system_idx, len(system_idx))
        acted_system_dim = [self.system_dim[idx] for idx in system_idx]
        self.append(KrausRepr(kraus_oper, system_idx, acted_system_dim))

    def stinespring_channel(self, stinespring_repr: Iterable[torch.Tensor],
                             system_idx: Union[Iterable[Iterable[int]], Iterable[int], int]) -> None:
        r"""Add custom channels in the Stinespring representation.

        Args:
            stinespring_repr: Stinespring representation.
            system_idx: Systems to apply the channel on.

        Examples:
            .. code-block:: python

                qc = Circuit(2)
                qc.stinespring_channel(stinespring_repr=eye(2), system_idx=[0])
                print(f'The latex code of this circuit is:\n{qc.to_latex()}')

            ::

                The latex code of this circuit is:
                \lstick{} & \gate[1]{\mathcal{N}}
        """
        system_idx = [system_idx] if isinstance(system_idx, int) else system_idx
        system_idx = self.register_idx(system_idx, len(system_idx))
        acted_system_dim = [self.system_dim[idx] for idx in system_idx]
        self.append(StinespringRepr(stinespring_repr, system_idx, acted_system_dim))
