"""
Quantum circuit: the state simulator.
Supports qubit (dim=2) and qudit (3 <= dim <= 36) systems.
For string-encoded samples/counts, digits use 0-9A-Z where A=10, ..., Z=35.
"""

# pylint: disable=invalid-name

from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
from functools import reduce
from operator import add

import numpy as np
import tensornetwork as tn

from . import gates
from . import channels
from .cons import backend, contractor, dtypestr, npdtype, _ALPHABET
from .quantum import QuOperator, identity, _infer_num_sites, onehot_d_tensor
from .simplify import _full_light_cone_cancel
from .basecircuit import BaseCircuit

Gate = gates.Gate
Tensor = Any


class Circuit(BaseCircuit):
    r"""
    ``Circuit`` class.
    Simple usage demo below.

    .. code-block:: python

        c = tc.Circuit(3)
        c.H(1)
        c.CNOT(0, 1)
        c.RX(2, theta=tc.num_to_tensor(1.))
        c.expectation([tc.gates.z(), (2, )]) # 0.54

    """

    is_dm = False

    def __init__(
        self,
        nqubits: int,
        inputs: Optional[Tensor] = None,
        mps_inputs: Optional[QuOperator] = None,
        tensors: Optional[Sequence[Tensor]] = None,
        split: Optional[Dict[str, Any]] = None,
        dim: Optional[int] = None,
    ) -> None:
        r"""
        Circuit object based on state simulator.
        Do not use this class with d!=2 directly, use tc.QuditCircuit instead for qudit systems.

        :param nqubits: The number of qubits in the circuit.
        :type nqubits: int
        :param dim: The local Hilbert space dimension per site. Qudit is supported for 2 <= d <= 36.
        :type dim: If None, the dimension of the circuit will be `2`, which is a qubit system.
        :param inputs: If not None, the initial state of the circuit is taken as ``inputs``
            instead of :math:`\vert 0 \rangle^n` qubits, defaults to None.
        :type inputs: Optional[Tensor], optional
        :param mps_inputs: QuVector for a MPS like initial wavefunction.
        :type mps_inputs: Optional[QuOperator]
        :param tensors: Sequence of tensors for a MPS like initial wavefunction.
            The order of legs for each tensor is assumed to be (bond-left, physical, bond-right).
        :type tensors: Optional[Sequence[Tensor]]
        :param split: dict if two qubit gate is ready for split, including parameters for at least one of
            ``max_singular_values`` and ``max_truncation_err``.
        :type split: Optional[Dict[str, Any]]
        """
        self._d = 2 if dim is None else dim
        self.inputs = inputs
        self.mps_inputs = mps_inputs
        self.split = split
        self._nqubits = nqubits

        self.circuit_param = {
            "nqubits": nqubits,
            "inputs": inputs,
            "mps_inputs": mps_inputs,
            "tensors": tensors,
            "split": split,
            "dim": dim,
        }
        if (inputs is None) and (mps_inputs is None) and (tensors is None):
            nodes = self.all_zero_nodes(nqubits, dim=self._d)
            self._front = [n.get_edge(0) for n in nodes]
        elif inputs is not None:  # provide input function
            inputs = backend.convert_to_tensor(inputs)
            inputs = backend.cast(inputs, dtype=dtypestr)
            inputs = backend.reshape(inputs, [-1])
            N = inputs.shape[0]
            n = _infer_num_sites(N, dim=self._d)
            assert n == nqubits or n == 2 * nqubits
            inputs = backend.reshape(inputs, [self._d] * n)
            inputs = Gate(inputs)
            nodes = [inputs]
            self._front = [inputs.get_edge(i) for i in range(n)]
        elif mps_inputs is not None:
            mps_nodes = list(mps_inputs.nodes)  # type: ignore
            for i, n in enumerate(mps_nodes):
                mps_nodes[i].tensor = backend.cast(n.tensor, dtypestr)  # type: ignore
            mps_edges = mps_inputs.out_edges + mps_inputs.in_edges  # type: ignore
            ndict, edict = tn.copy(mps_nodes)
            new_nodes = []
            for n in mps_nodes:
                new_nodes.append(ndict[n])
            new_front = []
            for e in mps_edges:
                new_front.append(edict[e])
            nodes = new_nodes
            self._front = new_front
        elif tensors is not None:
            nodes, self._front = self._tensors_to_nodes(tensors)
        else:
            raise ValueError("No inputs provided")  # should not be reached

        self.coloring_nodes(nodes, flag="inputs")
        self._nodes = nodes

        self._start_index = len(nodes)
        # self._start = nodes
        # self._meta_apply()

        # self._qcode = ""  # deprecated
        # self._qcode += str(self._nqubits) + "\n"
        self._qir: List[Dict[str, Any]] = []
        self._extra_qir: List[Dict[str, Any]] = []
        self._measure_counter = 0

    def replace_mps_inputs(self, mps_inputs: QuOperator) -> None:
        """
        Replace the input state in MPS representation while keep the circuit structure unchanged.

        :Example:
        >>> c = tc.Circuit(2)
        >>> c.X(0)
        >>>
        >>> c2 = tc.Circuit(2, mps_inputs=c.quvector())
        >>> c2.X(0)
        >>> c2.wavefunction()
        array([1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j], dtype=complex64)
        >>>
        >>> c3 = tc.Circuit(2)
        >>> c3.X(0)
        >>> c3.replace_mps_inputs(c.quvector())
        >>> c3.wavefunction()
        array([1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j], dtype=complex64)

        :param mps_inputs: (Nodes, dangling Edges) for a MPS like initial wavefunction.
        :type mps_inputs: Tuple[Sequence[Gate], Sequence[Edge]]
        """
        mps_nodes = mps_inputs.nodes
        mps_edges = mps_inputs.out_edges + mps_inputs.in_edges
        ndict, edict = tn.copy(mps_nodes)
        new_nodes = []
        for n in mps_nodes:
            new_nodes.append(ndict[n])
        new_front = []
        for e in mps_edges:
            new_front.append(edict[e])
        old = set(id(n) for n in self._nodes[: self._start_index])
        j = -1
        for n in self._nodes[: self._start_index]:
            for e in n:
                if e.is_dangling():
                    j += 1
                    self._front[j] = new_front[j]
                else:
                    if (id(e.node1) in old) and (id(e.node2) in old):
                        pass
                    else:
                        j += 1
                        if id(e.node2) == id(n):
                            other = (e.node1, e.axis1)
                        else:  # id(e.node1) == id(n):
                            other = (e.node2, e.axis2)
                        e.disconnect()
                        new_front[j] ^ other[0][other[1]]
        j += 1
        self._front += new_front[j:]
        self.coloring_nodes(new_nodes, flag="inputs")
        self._nodes = new_nodes + self._nodes[self._start_index :]
        self._start_index = len(new_nodes)

    # TODO(@refraction-ray): add noise support in IR
    # TODO(@refraction-ray): unify mid measure to basecircuit

    def _new_detector_work_circuit(self) -> BaseCircuit:
        from .densitymatrix import DMCircuit2

        return DMCircuit2(self._nqubits)

    def mid_measurement(self, index: int, keep: int = 0) -> Tensor:
        """
        Middle measurement in z-basis on the circuit, note the wavefunction output is not normalized
        with ``mid_measurement`` involved, one should normalize the state manually if needed.
        This is a post-selection method as keep is provided as a prior.

        :param index: The index of qubit that the Z direction postselection applied on.
        :type index: int
        :param keep: the post-selected digit in {0, ..., d-1}, defaults to be 0.
        :type keep: int, optional
        """
        # normalization not guaranteed
        gate = np.array(
            [[0.0] if _idx != keep else [1.0] for _idx in range(self._d)],
            dtype=npdtype,
        )

        mg1 = tn.Node(gate)
        mg2 = tn.Node(gate)
        # mg1.flag = "post-select"
        # mg1.is_dagger = False
        # mg1.id = id(mg1)
        # mg2.flag = "post-select"
        # mg2.is_dagger = False
        # mg2.id = id(mg2)
        self.coloring_nodes([mg1, mg2], flag="post-select")
        mg1.get_edge(0) ^ self._front[index]
        mg1.get_edge(1) ^ mg2.get_edge(1)
        self._front[index] = mg2.get_edge(0)
        self._nodes.append(mg1)
        self._nodes.append(mg2)
        r = backend.convert_to_tensor(keep)
        r = backend.cast(r, "int32")
        return r

    mid_measure = mid_measurement
    post_select = mid_measurement
    post_selection = mid_measurement

    def depolarizing2(
        self,
        index: int,
        *,
        px: float,
        py: float,
        pz: float,
        status: Optional[float] = None,
    ) -> float:
        """
        Apply a depolarizing channel to the circuit in a Monte Carlo way.
        For each call, one of the Pauli gates (X, Y, Z) or an Identity gate is applied to the qubit
        at the given index based on the probabilities `px`, `py`, and `pz`.

        :param index: The index of the qubit to apply the depolarizing channel on.
        :type index: int
        :param px: The probability of applying an X gate.
        :type px: float
        :param py: The probability of applying a Y gate.
        :type py: float
        :param pz: The probability of applying a Z gate.
        :type pz: float
        :param status: A random number between 0 and 1 to determine which gate to apply. If None,
                       a random number is generated automatically. Defaults to None.
        :type status: Optional[float], optional
        :return: Returns 0.0. The function modifies the circuit in place.
        :rtype: float
        """
        if status is None:
            status = backend.implicit_randu()[0]
        g = backend.cond(
            status < px,
            lambda: gates.x().tensor,  # type: ignore
            lambda: backend.cond(
                status < px + py,  # type: ignore
                lambda: gates.y().tensor,  # type: ignore
                lambda: backend.cond(
                    status < px + py + pz,  # type: ignore
                    lambda: gates.z().tensor,  # type: ignore
                    lambda: gates.i().tensor,  # type: ignore
                ),
            ),
        )
        # after implementing this, I realized that plain if is enough here for jit
        # the failure for previous implementation is because we use self.X(i) inside ``if``,
        # which has list append and incur bug in tensorflow jit
        # in terms of jax jit, the only choice is jax.lax.cond, since ``if tensor``` paradigm
        # is not supported in jax jit at all. (``Concrete Tensor Error``)
        self.any(index, unitary=g)  # type: ignore
        return 0.0
        # roughly benchmark shows that performance of two depolarizing in terms of
        # building time and running time are similar

    # overwritten now, deprecated
    def depolarizing_reference(
        self,
        index: int,
        *,
        px: float,
        py: float,
        pz: float,
        status: Optional[float] = None,
    ) -> Tensor:
        """
        Apply depolarizing channel in a Monte Carlo way,
        i.e. for each call of this method, one of gates from
        X, Y, Z, I are applied on the circuit based on the probability
        indicated by ``px``, ``py``, ``pz``.

        :param index: The qubit that depolarizing channel is on
        :type index: int
        :param px: probability for X noise
        :type px: float
        :param py: probability for Y noise
        :type py: float
        :param pz: probability for Z noise
        :type pz: float
        :param status: random seed uniformly from 0 to 1, defaults to None (generated implicitly)
        :type status: Optional[float], optional
        :return: int Tensor, the element lookup: [0: x, 1: y, 2: z, 3: I]
        :rtype: Tensor
        """

        # px/y/z here not support differentiation for now
        # jit compatible for now
        # assert px + py + pz < 1 and px >= 0 and py >= 0 and pz >= 0

        def step_function(x: Tensor) -> Tensor:
            r = (
                backend.sign(x - px)
                + backend.sign(x - px - py)
                + backend.sign(x - px - py - pz)
            )
            r = backend.cast(r / 2 + 1.5, dtype="int32")
            # [0: x, 1: y, 2: z, 3: I]

            return r

        if status is None:
            status = backend.implicit_randu()[0]
        r = step_function(status)
        rv = backend.onehot(r, 4)
        rv = backend.cast(rv, dtype=dtypestr)
        g = (
            rv[0] * gates.x().tensor  # type: ignore
            + rv[1] * gates.y().tensor  # type: ignore
            + rv[2] * gates.z().tensor  # type: ignore
            + rv[3] * gates.i().tensor  # type: ignore
        )
        self.any(index, unitary=g)  # type: ignore
        return r

    def unitary_kraus2(
        self,
        kraus: Sequence[Gate],
        *index: int,
        prob: Optional[Sequence[float]] = None,
        status: Optional[float] = None,
        name: Optional[str] = None,
    ) -> Tensor:
        """
        Apply a unitary Kraus channel to the circuit using a Monte Carlo approach. This method is functionally
        similar to `unitary_kraus` but uses `backend.switch` for selecting the Kraus operator, which can have
        different performance characteristics on some backends.

        A random Kraus operator from the provided list is applied to the circuit based on the given probabilities.
        This method is jittable and suitable for simulating noisy quantum circuits where the noise is represented
        by unitary Kraus operators.

        .. warning::
            This method may have issues with `vmap` due to potential concurrent access locks, potentially related with
            `backend.switch`. `unitary_kraus` is generally recommended.

        :param kraus: A sequence of `Gate` objects representing the unitary Kraus operators.
        :type kraus: Sequence[Gate]
        :param index: The qubit indices on which to apply the Kraus channel.
        :type index: int
        :param prob: A sequence of probabilities corresponding to each Kraus operator. If None, probabilities
                     are derived from the operators themselves. Defaults to None.
        :type prob: Optional[Sequence[float]], optional
        :param status: A random number between 0 and 1 to determine which Kraus operator to apply. If None,
                       a random number is generated automatically. Defaults to None.
        :type status: Optional[float], optional
        :param name: An optional name for the operation. Defaults to None.
        :type name: Optional[str], optional
        :return: A tensor indicating which Kraus operator was applied.
        :rtype: Tensor
        """

        # dont use, has issue conflicting with vmap, concurrent access lock emerged
        # potential issue raised from switch
        # general impl from Monte Carlo trajectory depolarizing above
        # still jittable
        # speed is similar to ``unitary_kraus``
        def index2gate2(r: Tensor, kraus: Sequence[Tensor]) -> Tensor:
            # r is int type Tensor of shape []
            return backend.switch(r, [lambda _=k: _ for k in kraus])  # type: ignore

        return self._unitary_kraus_template(
            kraus,
            *index,
            prob=prob,
            status=status,
            get_gate_from_index=index2gate2,
            name=name,
        )

    def unitary_kraus(
        self,
        kraus: Sequence[Gate],
        *index: int,
        prob: Optional[Sequence[float]] = None,
        status: Optional[float] = None,
        name: Optional[str] = None,
    ) -> Tensor:
        """
        Apply unitary gates in ``kraus`` randomly based on corresponding ``prob``.
        If ``prob`` is ``None``, this is reduced to kraus channel language.

        :param kraus: List of ``tc.gates.Gate`` or just Tensors
        :type kraus: Sequence[Gate]
        :param prob: prob list with the same size as ``kraus``, defaults to None
        :type prob: Optional[Sequence[float]], optional
        :param status: random seed between 0 to 1, defaults to None
        :type status: Optional[float], optional
        :return: shape [] int dtype tensor indicates which kraus gate is actually applied
        :rtype: Tensor
        """
        # general impl from Monte Carlo trajectory depolarizing above
        # still jittable

        def index2gate(r: Tensor, kraus: Sequence[Tensor]) -> Tensor:
            # r is int type Tensor of shape []
            l = len(kraus)
            r = backend.onehot(r, l)
            r = backend.cast(r, dtype=dtypestr)
            return reduce(add, [r[i] * kraus[i] for i in range(l)])

        return self._unitary_kraus_template(
            kraus,
            *index,
            prob=prob,
            status=status,
            get_gate_from_index=index2gate,
            name=name,
        )

    def _unitary_kraus_template(
        self,
        kraus: Sequence[Gate],
        *index: int,
        prob: Optional[Sequence[float]] = None,
        status: Optional[float] = None,
        get_gate_from_index: Optional[
            Callable[[Tensor, Sequence[Tensor]], Tensor]
        ] = None,
        name: Optional[str] = None,
    ) -> Tensor:  # DRY
        sites = len(index)
        kraus = [k.tensor if isinstance(k, tn.Node) else k for k in kraus]
        kraus = [gates.array_to_tensor(k) for k in kraus]
        kraus = [backend.reshapem(k) for k in kraus]
        if prob is None:
            prob = [
                backend.real(backend.trace(backend.adjoint(k) @ k) / k.shape[0])
                for k in kraus
            ]
            kraus = [
                k
                / backend.cast(
                    backend.where(
                        p
                        > backend.cast(
                            backend.convert_to_tensor(0.0), p.dtype  # type: ignore
                        ),
                        backend.sqrt(p),
                        backend.cast(
                            backend.convert_to_tensor(1.0), p.dtype  # type: ignore
                        ),
                    ),
                    dtypestr,
                )
                for k, p in zip(kraus, prob)
            ]
        if not backend.is_tensor(prob):
            prob = backend.convert_to_tensor(prob)
        prob_cumsum = backend.cumsum(prob)
        l = int(prob.shape[0])  # type: ignore

        def step_function(x: Tensor) -> Tensor:
            if l == 1:
                r = backend.convert_to_tensor(0.0)
            else:
                r = backend.sum(
                    backend.stack(
                        [backend.sign(x - prob_cumsum[i]) for i in range(l - 1)]
                    )
                )
            r = backend.cast(r / 2.0 + (l - 1) / 2.0, dtype="int32")
            # [0: kraus[0], 1: kraus[1]...]
            return r

        if status is None:
            status = backend.implicit_randu()[0]
        status = backend.convert_to_tensor(status)
        status = backend.real(status)
        prob_cumsum = backend.cast(prob_cumsum, dtype=status.dtype)  # type: ignore
        r = step_function(status)
        if get_gate_from_index is None:
            raise ValueError("no `get_gate_from_index` implementation is provided")
        g = get_gate_from_index(r, kraus)
        g = backend.reshape(g, [self._d] * (sites * 2))
        self.any(*index, unitary=g, name=name, dim=self._d)  # type: ignore
        return r

    def _general_kraus_tf(
        self,
        kraus: Sequence[Gate],
        *index: int,
        status: Optional[float] = None,
    ) -> float:
        # the graph building time is frustratingly slow, several minutes
        # though running time is in terms of ms
        sites = len(index)
        kraus_tensor = [k.tensor for k in kraus]
        kraus_tensor_f = [lambda _=k: _ for k in kraus_tensor]
        # must return tensor instead of ``tn.Node`` for switch`

        def calculate_kraus_p(i: Tensor) -> Tensor:
            # i: Tensor as int of shape []
            newnodes, newfront = self._copy()  # TODO(@refraction-ray): support reuse?
            # simply reuse=True is wrong, as the circuit is contracting at building
            # self._copy seems slower than self._copy_state, but anyway the building time is unacceptable
            lnewnodes, lnewfront = self._copy(conj=True)
            kraus_i = backend.switch(i, kraus_tensor_f)
            k = gates.Gate(kraus_i)
            kc = gates.Gate(backend.conj(kraus_i))
            # begin connect
            for ind, j in enumerate(index):
                newfront[j] ^ k[ind + sites]
                k[ind] ^ kc[ind]
                kc[ind + sites] ^ lnewfront[j]
            for j in range(self._nqubits):
                if j not in index:
                    newfront[j] ^ lnewfront[j]
            norm_square = contractor(newnodes + lnewnodes + [k, kc]).tensor
            return backend.real(norm_square)

        if status is None:
            status = backend.implicit_randu()[0]

        import tensorflow as tf  # tf only implementation

        weight = 1.0
        fallback_weight = 0.0
        fallback_weight_i = 0
        len_kraus = len(kraus)
        for i in tf.range(len_kraus):  # breaks backend agnostic
            # nested for and if, if tensor inner must come with for in tensor outer, s.t. autograph works
            weight = calculate_kraus_p(i)
            if weight > fallback_weight:
                fallback_weight_i = i
                fallback_weight = weight
            status -= weight
            if status < 0:
                # concern here, correctness not sure in tf jit, fail anyway in jax jit
                break
        # placing a Tensor-dependent break, continue or return inside a Python loop
        # https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/common_errors.md

        if (
            status >= 0 or weight == 0
        ):  # the same concern, but this simple if is easy to convert to ``backend.cond``
            # Floating point error resulted in a malformed sample.
            # Fall back to the most likely case.
            # inspired from cirq implementation (Apcache 2).
            weight = fallback_weight
            i = fallback_weight_i
        kraus_i = backend.switch(i, kraus_tensor_f)
        newgate = kraus_i / backend.cast(backend.sqrt(weight), dtypestr)
        self.any(*index, unitary=newgate)  # type: ignore
        return 0.0

    def _general_kraus_2(
        self,
        kraus: Sequence[Gate],
        *index: int,
        status: Optional[float] = None,
        with_prob: bool = False,
        name: Optional[str] = None,
    ) -> Tensor:
        # the graph building time is frustratingly slow, several minutes
        # though running time is in terms of ms
        # raw running time in terms of s
        # note jax gpu building time is fast, in the order of 10s.!!
        # the typical scenario we are talking: 10 qubits, 3 layers of entangle gates and 3 layers of noise
        # building for jax+GPU ~100s 12 qubit * 5 layers
        # 370s 14 qubit * 7 layers, 0.35s running on vT4
        # vmap, grad, vvag are all fine for this function
        # layerwise jit technique can greatly boost the staging time, see in /examples/mcnoise_boost.py
        sites = len(index)
        kraus_tensor = [k.tensor if isinstance(k, tn.Node) else k for k in kraus]
        kraus_tensor = [gates.array_to_tensor(k) for k in kraus_tensor]

        # tn with hole
        newnodes, newfront = self._copy()
        lnewnodes, lnewfront = self._copy(conj=True)
        des = [newfront[j] for j in index] + [lnewfront[j] for j in index]
        for j in range(self._nqubits):
            if j not in index:
                newfront[j] ^ lnewfront[j]
        ns = contractor(newnodes + lnewnodes, output_edge_order=des)
        ntensor = ns.tensor
        # ns, des

        def calculate_kraus_p(i: int) -> Tensor:
            # i: Tensor as int of shape []
            # kraus_i = backend.switch(i, kraus_tensor_f)
            kraus_i = kraus_tensor[i]
            dm = gates.Gate(ntensor)
            k = gates.Gate(kraus_i)
            kc = gates.Gate(backend.conj(kraus_i))
            # begin connect
            for ind in range(sites):
                dm[ind] ^ k[ind + sites]
                k[ind] ^ kc[ind]
                kc[ind + sites] ^ dm[ind + sites]
            norm_square = contractor([dm, k, kc]).tensor
            return backend.real(norm_square)

        prob = [calculate_kraus_p(i) for i in range(len(kraus))]
        eps = 1e-10
        new_kraus = [
            k / backend.cast(backend.sqrt(w) + eps, dtypestr)
            for w, k in zip(prob, kraus_tensor)
        ]
        pick = self.unitary_kraus(
            new_kraus,
            *index,
            prob=prob,
            status=status,
            name=name,
        )
        if not with_prob:
            return pick
        else:
            return pick, prob

    def general_kraus(
        self,
        kraus: Sequence[Gate],
        *index: int,
        status: Optional[float] = None,
        with_prob: bool = False,
        name: Optional[str] = None,
    ) -> Tensor:
        """
        Monte Carlo trajectory simulation of general Kraus channel whose Kraus operators cannot be
        amplified to unitary operators. For unitary operators composed Kraus channel, :py:meth:`unitary_kraus`
        is much faster.

        This function is jittable in theory. But only jax+GPU combination is recommended for jit
        since the graph building time is too long for other backend options; though the running
        time of the function is very fast for every case.

        :param kraus: A list of ``tn.Node`` for Kraus operators.
        :type kraus: Sequence[Gate]
        :param index: The qubits index that Kraus channel is applied on.
        :type index: int
        :param status: Random tensor uniformly between 0 or 1, defaults to be None,
            when the random number will be generated automatically
        :type status: Optional[float], optional
        """
        return self._general_kraus_2(
            kraus,
            *index,
            status=status,
            with_prob=with_prob,
            name=name,
        )

    apply_general_kraus = general_kraus

    @staticmethod
    def apply_general_kraus_delayed(
        krausf: Callable[..., Sequence[Gate]], is_unitary: bool = False
    ) -> Callable[..., None]:
        def apply(
            self: "Circuit",
            *index: int,
            status: Optional[float] = None,
            name: Optional[str] = None,
            **vars: float,
        ) -> None:
            channel_parameters = dict(vars)
            kraus = krausf(**vars)
            if name is None:
                name = getattr(kraus, "name", None)
            if name is None:
                name = getattr(krausf, "__name__", "channel").replace("channel", "")
            if not is_unitary:
                self.apply_general_kraus(kraus, *index, status=status, name=name)
            else:
                self.unitary_kraus(kraus, *index, status=status, name=name)
            if len(self._qir) > 0:
                self._qir[-1]["name"] = name
                self._qir[-1]["is_channel"] = True
                self._qir[-1]["channel_f"] = krausf
                self._qir[-1]["channel_parameters"] = channel_parameters
                self._qir[-1]["channel_unitary"] = is_unitary

        return apply

    @classmethod
    def _meta_apply_channels(cls) -> None:
        for k in channels.channels:
            if k in ["depolarizing", "generaldepolarizing"]:
                is_unitary = True
            else:
                is_unitary = False
            setattr(
                cls,
                k,
                cls.apply_general_kraus_delayed(
                    getattr(channels, k + "channel"), is_unitary=is_unitary
                ),
            )
            doc = """
            Apply %s quantum channel on the circuit.
            See :py:meth:`tensorcircuit.channels.%schannel`

            :param index: Site index that the gate applies on.
            :type index: int.
            :param status: uniform external random number between 0 and 1
            :type status: Tensor
            :param vars: Parameters for the channel.
            :type vars: float.
            """ % (
                k,
                k,
            )
            getattr(cls, k).__doc__ = doc

    def is_valid(self) -> bool:
        """
        [WIP], check whether the circuit is legal.

        :return: The bool indicating whether the circuit is legal
        :rtype: bool
        """
        try:
            assert len(self._front) == self._nqubits
            for n in self._nodes:
                for e in n.get_all_dangling():
                    assert e in self._front
            return True
        except AssertionError:
            return False

    def wavefunction(self, form: str = "default") -> tn.Node.tensor:
        """
        Compute the output wavefunction from the circuit.

        :param form: The str indicating the form of the output wavefunction.
            "default": [-1], "ket": [-1, 1], "bra": [1, -1]
        :type form: str, optional
        :return: Tensor with the corresponding shape.
        :rtype: Tensor
        """
        nodes, d_edges = self._copy()
        t = contractor(nodes, output_edge_order=d_edges)
        if form == "default":
            shape = [-1]
        elif form == "ket":
            shape = [-1, 1]
        elif form == "bra":  # no conj here
            shape = [1, -1]
        return backend.reshape(t.tensor, shape=shape)

    state = wavefunction

    def get_quoperator(self) -> QuOperator:
        """
        Get the ``QuOperator`` MPO like representation of the circuit unitary without contraction.

        :return: ``QuOperator`` object for the circuit unitary (open indices for the input state)
        :rtype: QuOperator
        """
        mps = identity([self._d] * self._nqubits)
        c = Circuit(self._nqubits, dim=self._d)
        ns, es = self._copy()
        c._nodes = ns
        c._front = es
        c.replace_mps_inputs(mps)
        return QuOperator(c._front[: self._nqubits], c._front[self._nqubits :])

    quoperator = get_quoperator
    # both are not good names, but for backward compatibility

    get_circuit_as_quoperator = get_quoperator
    get_state_as_quvector = BaseCircuit.quvector

    def matrix(self) -> Tensor:
        """
        Get the unitary matrix for the circuit irrespective with the circuit input state.

        :return: The circuit unitary matrix
        :rtype: Tensor
        """
        mps = identity([self._d] * self._nqubits)
        c = Circuit(self._nqubits, dim=self._d)
        ns, es = self._copy()
        c._nodes = ns
        c._front = es
        c.replace_mps_inputs(mps)
        return backend.reshapem(c.state())

    def measure_reference(
        self, *index: int, with_prob: bool = False
    ) -> Tuple[str, float]:
        """
        Take measurement on the given quantum lines by ``index``.

        Return format:
        - For d <= 36, the sample is a base-d string using 0-9A-Z (A=10,...).

        :Example:

        >>> c = tc.Circuit(3)
        >>> c.H(0)
        >>> c.h(1)
        >>> c.toffoli(0, 1, 2)
        >>> c.measure(2)
        ('1', -1.0)
        >>> # Another possible output: ('0', -1.0)
        >>> c.measure(2, with_prob=True)
        ('1', (0.25000011920928955+0j))
        >>> # Another possible output: ('0', (0.7499998807907104+0j))

        :param index: Measure on which quantum line.
        :param with_prob: If true, theoretical probability is also returned.
        :return: The sample output and probability (optional) of the quantum line.
        :rtype: Tuple[str, float]
        """
        # not jit compatible due to random number generations!
        sample = ""
        p = 1.0
        for j in index:
            nodes1, edge1 = self._copy()
            nodes2, edge2 = self._copy(conj=True)
            for i, e in enumerate(edge1):
                if i != j:
                    e ^ edge2[i]
            for i in range(len(sample)):
                m = onehot_d_tensor(sample[i], d=self._d)
                nodes1.append(tn.Node(m))
                nodes1[-1].get_edge(0) ^ edge1[index[i]]
                nodes2.append(tn.Node(m))
                nodes2[-1].get_edge(0) ^ edge2[index[i]]
            nodes1.extend(nodes2)
            rho = (
                1
                / p
                * contractor(nodes1, output_edge_order=[edge1[j], edge2[j]]).tensor
            )
            probs = backend.real(backend.diagonal(rho))
            probs /= backend.sum(probs)
            outcome = backend.implicit_randc(self._d, shape=1, p=probs)

            sample += _ALPHABET[outcome]
            p *= float(probs[outcome])

        if with_prob:
            return sample, p
        else:
            return sample, -1.0

    # TODO(@refraction-ray): more _before function like state_before? and better API?

    def expectation(
        self,
        *ops: Tuple[tn.Node, List[int]],
        reuse: bool = True,
        enable_lightcone: bool = False,
        noise_conf: Optional[Any] = None,
        nmc: int = 1000,
        status: Optional[Tensor] = None,
        **kws: Any,
    ) -> Tensor:
        """
        Compute the expectation of corresponding operators.
        For qudit (d > 2),
        ensure that operator tensor shapes are consistent with d (each site contributes two axes of size d).

        Noise shorthand (via noise_conf) is qubit-only; for d>2, use explicit operators.

        :Example:

        >>> c = tc.Circuit(2)
        >>> c.H(0)
        >>> c.expectation((tc.gates.z(), [0]))
        array(0.+0.j, dtype=complex64)

        >>> c = tc.Circuit(2)
        >>> c.cnot(0, 1)
        >>> c.rx(0, theta=0.4)
        >>> c.rx(1, theta=0.8)
        >>> c.h(0)
        >>> c.h(1)
        >>> error1 = tc.channels.generaldepolarizingchannel(0.1, 1)
        >>> error2 = tc.channels.generaldepolarizingchannel(0.06, 2)
        >>> noise_conf = NoiseConf()
        >>> noise_conf.add_noise("rx", error1)
        >>> noise_conf.add_noise("cnot", [error2], [[0, 1]])
        >>> c.expectation((tc.gates.x(), [0]), noise_conf=noise_conf, nmc=10000)
        (0.46274087-3.764033e-09j)

        :param ops: Operator and its position on the circuit,
            eg. ``(tc.gates.z(), [1, ]), (tc.gates.x(), [2, ])`` is for operator :math:`Z_1X_2`.
        :type ops: Tuple[tn.Node, List[int]]
        :param reuse: If True, then the wavefunction tensor is cached for further expectation evaluation,
            defaults to be true.
        :type reuse: bool, optional
        :param enable_lightcone: whether enable light cone simplification, defaults to False
        :type enable_lightcone: bool, optional
        :param noise_conf: Noise Configuration, defaults to None
        :type noise_conf: Optional[NoiseConf], optional
        :param nmc: repetition time for Monte Carlo sampling for noisfy calculation, defaults to 1000
        :type nmc: int, optional
        :param status: external randomness given by tensor uniformly from [0, 1], defaults to None,
            used for noisy circuit sampling
        :type status: Optional[Tensor], optional
        :raises ValueError: "Cannot measure two operators in one index"
        :return: Tensor with one element
        :rtype: Tensor
        """
        if noise_conf is None:
            # if not reuse:
            #     nodes1, edge1 = self._copy()
            #     nodes2, edge2 = self._copy(conj=True)
            # else:  # reuse

            # self._nodes = nodes1
            if enable_lightcone:
                reuse = False
            nodes1 = self.expectation_before(*ops, reuse=reuse)
            if enable_lightcone:
                nodes1 = _full_light_cone_cancel(nodes1)
            return contractor(nodes1).tensor
        else:
            from .noisemodel import expectation_noisfy

            return expectation_noisfy(
                self,
                *ops,
                noise_conf=noise_conf,
                nmc=nmc,
                status=status,
                **kws,
            )


Circuit._meta_apply()
Circuit._meta_apply_channels()


def expectation(
    *ops: Tuple[tn.Node, List[int]],
    ket: Tensor,
    bra: Optional[Tensor] = None,
    conj: bool = True,
    normalization: bool = False,
    dim: Optional[int] = None,
) -> Tensor:
    """
    Compute :math:`\\langle bra\\vert ops \\vert ket\\rangle`.
    For qudit systems (d>2), ops must be reshaped with per-site axes of length d.

    Example 1 (:math:`bra` is same as :math:`ket`)

    >>> c = tc.Circuit(3)
    >>> c.H(0)
    >>> c.ry(1, theta=tc.num_to_tensor(0.8 + 0.7j))
    >>> c.cnot(1, 2)
    >>> state = c.wavefunction() # the state of this circuit
    >>> x1z2 = [(tc.gates.x(), [0]), (tc.gates.z(), [1])] # input qubits
    >>>
    >>> # Expection of this circuit / <state|*x1z2|state>
    >>> c.expectation(*x1z2)
    array(0.69670665+0.j, dtype=complex64)
    >>> tc.expectation(*x1z2, ket=state)
    (0.6967066526412964+0j)
    >>>
    >>> # Normalize(expection of Circuit) / Normalize(<state|*x1z2|state>)
    >>> c.expectation(*x1z2) / tc.backend.norm(state) ** 2
    (0.5550700389340034+0j)
    >>> tc.expectation(*x1z2, ket=state, normalization=True)
    (0.55507004+0j)

    Example 2 (:math:`bra` is different from :math:`ket`)

    >>> c = tc.Circuit(2)
    >>> c.X(1)
    >>> s1 = c.state()
    >>> c2 = tc.Circuit(2)
    >>> c2.X(0)
    >>> s2 = c2.state()
    >>> c3 = tc.Circuit(2)
    >>> c3.H(1)
    >>> s3 = c3.state()
    >>> x1x2 = [(tc.gates.x(), [0]), (tc.gates.x(), [1])]
    >>>
    >>> tc.expectation(*x1x2, ket=s1, bra=s2)
    (1+0j)
    >>> tc.expectation(*x1x2, ket=s3, bra=s2)
    (0.7071067690849304+0j) # 1/sqrt(2)

    :param ket: :math:`ket`. The state in tensor or ``QuVector`` format
    :type ket: Tensor
    :param bra: :math:`bra`, defaults to None, which is the same as ``ket``.
    :type bra: Optional[Tensor], optional
    :param dim: dimension of the circuit (defaults to 2)
    :type dim: int, optional
    :param conj: :math:`bra` changes to the adjoint matrix of :math:`bra`, defaults to True.
    :type conj: bool, optional
    :param normalization: Normalize the :math:`ket` and :math:`bra`, defaults to False.
    :type normalization: bool, optional
    :raises ValueError: "Cannot measure two operators in one index"
    :return: The result of :math:`\\langle bra\\vert ops \\vert ket\\rangle`.
    :rtype: Tensor
    """
    dim = 2 if dim is None else dim
    if bra is None:
        bra = ket
    if isinstance(ket, QuOperator):
        if conj is True:
            bra = bra.adjoint()
        # TODO(@refraction-ray) omit normalization arg for now
        n = len(ket.out_edges)
        occupied = set()
        nodes = list(ket.nodes) + list(bra.nodes)
        # TODO(@refraction-ray): is the order guaranteed or affect some types of contractor?
        for op, index in ops:
            if not isinstance(op, tn.Node):
                # op is only a matrix
                op = backend.reshaped(op, dim)
                op = gates.Gate(op)
            if isinstance(index, int):
                index = [index]
            noe = len(index)
            for j, e in enumerate(index):
                if e in occupied:
                    raise ValueError("Cannot measure two operators in one index")
                bra.in_edges[e] ^ op.get_edge(j)
                ket.out_edges[e] ^ op.get_edge(j + noe)
                occupied.add(e)
            nodes.append(op)
        for j in range(n):
            if j not in occupied:  # edge1[j].is_dangling invalid here!
                ket.out_edges[j] ^ bra.in_edges[j]
        # self._nodes = nodes1
        num = contractor(nodes).tensor
        return num

    else:
        # ket is the tensor
        if conj is True:
            bra = backend.conj(bra)
        ket = backend.reshape(ket, [-1])
        ket = backend.reshaped(ket, dim)
        bra = backend.reshaped(bra, dim)
        n = len(backend.shape_tuple(ket))
        ket = Gate(ket)
        bra = Gate(bra)
        occupied = set()
        nodes = [ket, bra]
        if normalization is True:
            normket = backend.norm(ket.tensor)
            normbra = backend.norm(bra.tensor)
        for op, index in ops:
            if not isinstance(op, tn.Node):
                # op is only a matrix
                op = backend.reshaped(op, dim)
                op = gates.Gate(op)
            if isinstance(index, int):
                index = [index]
            noe = len(index)
            for j, e in enumerate(index):
                if e in occupied:
                    raise ValueError("Cannot measure two operators in one index")
                bra[e] ^ op.get_edge(j)
                ket[e] ^ op.get_edge(j + noe)
                occupied.add(e)
            nodes.append(op)
        for j in range(n):
            if j not in occupied:  # edge1[j].is_dangling invalid here!
                ket[j] ^ bra[j]
        # self._nodes = nodes1
        num = contractor(nodes).tensor
        if normalization is True:
            den = normket * normbra
        else:
            den = 1.0
        return num / den
