"""
Constants and setups
"""

# pylint: disable=invalid-name

import logging
import sys
import time
from contextlib import contextmanager
from functools import partial, reduce, wraps, lru_cache
from operator import mul
from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple

import numpy as np
import opt_einsum
import tensornetwork as tn
from tensornetwork.backend_contextmanager import get_default_backend
from networkx.utils import UnionFind

from .backends.numpy_backend import NumpyBackend
from .backends import get_backend
from .simplify import _multi_remove

logger = logging.getLogger(__name__)

## monkey patch
_NODE_CREATION_COUNTER = 0
_original_node_init = tn.Node.__init__
_original_copynode_init = tn.CopyNode.__init__


@wraps(_original_node_init)
def _patched_node_init(self: Any, *args: Any, **kwargs: Any) -> None:
    """Patched Node.__init__ to add a stable creation ID."""
    global _NODE_CREATION_COUNTER
    _original_node_init(self, *args, **kwargs)
    self._stable_id_ = _NODE_CREATION_COUNTER
    _NODE_CREATION_COUNTER += 1


@wraps(_original_copynode_init)
def _patched_copynode_init(self: Any, *args: Any, **kwargs: Any) -> None:
    """Patched CopyNode.__init__ to add a stable creation ID."""
    global _NODE_CREATION_COUNTER
    _original_copynode_init(self, *args, **kwargs)
    self._stable_id_ = _NODE_CREATION_COUNTER
    _NODE_CREATION_COUNTER += 1


tn.Node.__init__ = _patched_node_init
tn.CopyNode.__init__ = _patched_copynode_init


def _get_edge_stable_key(edge: tn.Edge) -> Tuple[int, int, int, int]:
    n1, n2 = edge.node1, edge.node2
    id1 = getattr(n1, "_stable_id_", -1)
    id2 = getattr(n2, "_stable_id_", -1) if n2 is not None else -2  # -2 for dangling

    if id1 > id2 or (id1 == id2 and edge.axis1 > edge.axis2):
        id1, id2, ax1, ax2 = id2, id1, edge.axis2, edge.axis1
    else:
        ax1, ax2 = edge.axis1, edge.axis2
    return (id1, ax1, id2, ax2)


def sorted_edges(edges: Iterator[tn.Edge]) -> List[tn.Edge]:
    return sorted(edges, key=_get_edge_stable_key)


package_name = "tensorcircuit"
thismodule = sys.modules[__name__]
dtypestr = "complex64"
rdtypestr = "float32"
idtypestr = "int32"
npdtype = np.complex64
backend: NumpyBackend = get_backend("numpy")
contractor = tn.contractors.auto
_ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
# these above lines are just for mypy, it is not very good at evaluating runtime object


def _set_global_contractor(contractor_fn: Callable[..., Any]) -> None:
    for module in sys.modules:
        if module.startswith(package_name):
            setattr(sys.modules[module], "contractor", contractor_fn)


def set_tensornetwork_backend(
    backend: Optional[str] = None, set_global: bool = True
) -> Any:
    r"""To set the runtime backend of tensorcircuit.

    Note: ``tc.set_backend`` and ``tc.cons.set_tensornetwork_backend`` are the same.

    :Example:

    >>> tc.set_backend("numpy")
    numpy_backend
    >>> tc.gates.num_to_tensor(0.1)
    array(0.1+0.j, dtype=complex64)
    >>>
    >>> tc.set_backend("tensorflow")
    tensorflow_backend
    >>> tc.gates.num_to_tensor(0.1)
    <tf.Tensor: shape=(), dtype=complex64, numpy=(0.1+0j)>
    >>>
    >>> tc.set_backend("pytorch")
    pytorch_backend
    >>> tc.gates.num_to_tensor(0.1)
    tensor(0.1000+0.j)
    >>>
    >>> tc.set_backend("jax")
    jax_backend
    >>> tc.gates.num_to_tensor(0.1)
    DeviceArray(0.1+0.j, dtype=complex64)

    :param backend: "numpy", "tensorflow", "jax", "pytorch". defaults to None,
        which gives the same behavior as ``tensornetwork.backend_contextmanager.get_default_backend()``.
    :type backend: Optional[str], optional
    :param set_global: Whether the object should be set as global.
    :type set_global: bool
    :return: The `tc.backend` object that with all registered universal functions.
    :rtype: backend object
    """
    if not backend:
        backend = get_default_backend()
    backend_obj = get_backend(backend)
    if set_global:
        for module in sys.modules:
            if module.startswith(package_name):
                setattr(sys.modules[module], "backend", backend_obj)
        tn.set_default_backend(backend)
    return backend_obj


set_backend = set_tensornetwork_backend

set_tensornetwork_backend()


def set_function_backend(backend: Optional[str] = None) -> Callable[..., Any]:
    """
    Function decorator to set function-level runtime backend

    :param backend: "numpy", "tensorflow", "jax", "pytorch", defaults to None
    :type backend: Optional[str], optional
    :return: Decorated function
    :rtype: Callable[..., Any]
    """

    def wrapper(f: Callable[..., Any]) -> Callable[..., Any]:
        @wraps(f)
        def newf(*args: Any, **kws: Any) -> Any:
            old_backend = getattr(thismodule, "backend").name
            set_backend(backend)
            try:
                r = f(*args, **kws)
            finally:
                set_backend(old_backend)
            return r

        return newf

    return wrapper


@contextmanager
def runtime_backend(backend: Optional[str] = None) -> Iterator[Any]:
    """
    Context manager to set with-level runtime backend

    :param backend: "numpy", "tensorflow", "jax", "pytorch", defaults to None
    :type backend: Optional[str], optional
    :yield: the backend object
    :rtype: Iterator[Any]
    """
    old_backend = getattr(thismodule, "backend").name
    K = set_backend(backend)
    yield K
    set_backend(old_backend)


def set_dtype(dtype: Optional[str] = None, set_global: bool = True) -> Tuple[str, str]:
    """
    Set the global runtime numerical dtype of tensors.

    :param dtype: "complex64"/"float32" or "complex128"/"float64",
        defaults to None, which is equivalent to "complex64".
    :type dtype: Optional[str], optional
    :return: complex dtype str and the corresponding real dtype str
    :rtype: Tuple[str, str]
    """
    if not dtype:
        dtype = "complex64"

    if dtype == "complex64":
        rdtype = "float32"
    elif dtype == "complex128":
        rdtype = "float64"
    elif dtype == "float32":
        dtype = "complex64"
        rdtype = "float32"
    elif dtype == "float64":
        dtype = "complex128"
        rdtype = "float64"
    else:
        raise ValueError(f"Unsupported data type: {dtype}")

    if dtype == "complex128":
        idtype = "int64"
    else:
        idtype = "int32"

    try:
        from jax import config
    except ImportError:
        config = None  # type: ignore

    if config is not None:
        if dtype == "complex128":
            config.update("jax_enable_x64", True)
        elif dtype == "complex64":
            config.update("jax_enable_x64", False)

    if set_global:
        npdtype = getattr(np, dtype)
        for module in sys.modules:
            if module.startswith(package_name):
                setattr(sys.modules[module], "dtypestr", dtype)
                setattr(sys.modules[module], "rdtypestr", rdtype)
                setattr(sys.modules[module], "idtypestr", idtype)
                setattr(sys.modules[module], "npdtype", npdtype)

        from .gates import meta_gate

        meta_gate()
    return dtype, rdtype


get_dtype = partial(set_dtype, set_global=False)

set_dtype()


def set_function_dtype(dtype: Optional[str] = None) -> Callable[..., Any]:
    """
    Function decorator to set function-level numerical dtype

    :param dtype: "complex64" or "complex128", defaults to None
    :type dtype: Optional[str], optional
    :return: The decorated function
    :rtype: Callable[..., Any]
    """

    def wrapper(f: Callable[..., Any]) -> Callable[..., Any]:
        @wraps(f)
        def newf(*args: Any, **kws: Any) -> Any:
            old_dtype = getattr(thismodule, "dtypestr")
            set_dtype(dtype)
            r = f(*args, **kws)
            set_dtype(old_dtype)
            return r

        return newf

    return wrapper


@contextmanager
def runtime_dtype(dtype: Optional[str] = None) -> Iterator[Tuple[str, str]]:
    """
    Context manager to set with-level runtime dtype

    :param dtype: "complex64" or "complex128", defaults to None ("complex64")
    :type dtype: Optional[str], optional
    :yield: complex dtype str and real dtype str
    :rtype: Iterator[Tuple[str, str]]
    """
    old_dtype = getattr(thismodule, "dtypestr")
    dtuple = set_dtype(dtype)
    yield dtuple
    set_dtype(old_dtype)


# here below comes other contractors (just works,
# but correctness has not been extensively tested for some of them)


def _sizen(node: tn.Node, is_log: bool = False) -> int:
    s = reduce(mul, node.tensor.shape + (1,))
    if is_log:
        return int(np.log2(s))
    return s  # type: ignore


def _merge_single_gates(
    nodes: List[Any], total_size: Optional[int] = None
) -> Tuple[List[Any], int]:
    # TODO(@refraction-ray): investigate whether too much copy here so that staging is slow for large circuit
    nodes = list(nodes)
    if total_size is None:
        total_size = sum([_sizen(t) for t in nodes])
    queue = [n for n in nodes if len(n.tensor.shape) <= 2]
    while queue:
        n0 = queue[0]
        try:
            n0[0]
        except IndexError:
            queue = _multi_remove(queue, [0])
            continue
        if n0[0].is_dangling():
            try:
                e0 = n0[1]
                if e0.is_dangling():
                    queue = _multi_remove(queue, [0])
                    continue
            except IndexError:
                queue = _multi_remove(queue, [0])
                continue
        else:
            e0 = n0[0]
        njs = [i for i, n in enumerate(nodes) if id(n) in [id(e0.node1), id(e0.node2)]]
        qjs = [i for i, n in enumerate(queue) if id(n) in [id(e0.node1), id(e0.node2)]]
        new_node = tn.contract_parallel(e0)
        total_size += _sizen(new_node)

        logger.debug(
            _sizen(new_node, is_log=True),
        )
        queue = _multi_remove(queue, qjs)
        if len(njs) > 1:
            nodes[njs[1]] = new_node
            nodes = _multi_remove(nodes, [njs[0]])
        else:  # trace edge?
            nodes[njs[0]] = new_node
        if len(new_node.tensor.shape) <= 2:
            # queue.append(new_node)
            queue.insert(0, new_node)
    return nodes, total_size


def experimental_contractor(
    nodes: List[Any],
    output_edge_order: Optional[List[Any]] = None,
    ignore_edge_order: bool = False,
    local_steps: int = 2,
) -> Any:
    total_size = sum([_sizen(t) for t in nodes])
    nodes = list(nodes)
    # merge single qubit gate
    if len(nodes) > 5:
        nodes, total_size = _merge_single_gates(nodes, total_size)

    # further dq fusion
    if len(nodes) > 15:
        for r in range(local_steps):
            if len(nodes) < 10:
                break
            i = 0
            while len(nodes) > i + 1:
                new_node = tn.contract_between(
                    nodes[i], nodes[i + 1], allow_outer_product=True
                )
                total_size += _sizen(new_node)

                logger.debug(
                    r,
                    _sizen(new_node, is_log=True),
                )
                nodes[i] = new_node
                nodes = _multi_remove(nodes, [i + 1])
                i += 1

    logger.info("length of remaining nodes after dq fusion: %s" % len(nodes))
    nodes = list(reversed(nodes))

    while len(nodes) > 1:
        new_node = tn.contract_between(nodes[-1], nodes[-2], allow_outer_product=True)
        nodes = _multi_remove(nodes, [len(nodes) - 2, len(nodes) - 1])
        nodes.append(new_node)
        logger.debug(_sizen(new_node, is_log=True))
        total_size += _sizen(new_node)

    logger.info("----- WRITE: %s --------\n" % np.log2(total_size))

    # if the final node has more than one edge,
    # output_edge_order has to be specified
    final_node = nodes[0]
    if output_edge_order is not None:
        final_node.reorder_edges(output_edge_order)
    return final_node


def plain_contractor(
    nodes: List[Any],
    output_edge_order: Optional[List[Any]] = None,
    ignore_edge_order: bool = False,
) -> Any:
    """
    The naive state-vector simulator contraction path.

    :param nodes: The list of ``tn.Node``.
    :type nodes: List[Any]
    :param output_edge_order: The list of dangling node edges, defaults to be None.
    :type output_edge_order: Optional[List[Any]], optional
    :return: The ``tn.Node`` after contraction
    :rtype: tn.Node
    """
    total_size = sum([_sizen(t) for t in nodes])
    # nodes = list(reversed(list(nodes)))
    nodes = list(nodes)
    nodes = list(reversed(nodes))

    width = 0

    while len(nodes) > 1:
        new_node = tn.contract_between(nodes[-1], nodes[-2], allow_outer_product=True)
        nodes = _multi_remove(nodes, [len(nodes) - 2, len(nodes) - 1])
        nodes.append(new_node)
        im_size = _sizen(new_node, is_log=True)
        logger.debug(im_size)
        width = max(width, im_size)
        total_size += _sizen(new_node)
    logger.info("----- SIZE: %s --------\n" % width)
    logger.info("----- WRITE: %s --------\n" % np.log2(total_size))

    final_node = nodes[0]
    if output_edge_order is not None:
        final_node.reorder_edges(output_edge_order)
    return final_node


# TODO(@refraction-ray): consistent logger system for different contractors.


def nodes_to_adj(ns: List[Any]) -> Any:
    ind = {id(n): i for i, n in enumerate(ns)}
    adj = np.zeros([len(ns), len(ns)])
    for node in ns:
        for e in node:
            if not e.is_dangling():
                if id(e.node1) == id(node):
                    onode = e.node2
                else:
                    onode = e.node1
                adj[ind[id(node)], ind[id(onode)]] += np.log10(e.dimension)
    return adj


try:
    _ps = tn.contractors.custom_path_solvers.pathsolvers
    has_ps = True
except AttributeError:
    has_ps = False


def d2s(n: int, dl: List[Any]) -> List[Any]:
    # dynamic to static list
    nums = [i for i in range(n)]
    i = n
    sl = []
    for a, b in dl:
        sl.append([nums[a], nums[b]])
        nums = _multi_remove(nums, [a, b])
        nums.insert(a, i)
        i += 1
    return sl


# seems worse than plain contraction in most cases
def tn_greedy_contractor(
    nodes: List[Any],
    output_edge_order: Optional[List[Any]] = None,
    ignore_edge_order: bool = False,
    max_branch: int = 1,
) -> Any:
    nodes = list(nodes)
    adj = nodes_to_adj(nodes)
    path = _ps.full_solve_complete(adj, max_branch=max_branch)[0]
    dl = []
    for i in range(path.shape[1]):
        a, b = path[:, i]
        dl.append([a, b])
    sl = d2s(len(nodes), dl)
    for a, b in sl:
        new_node = tn.contract_between(nodes[a], nodes[b], allow_outer_product=True)
        nodes.append(new_node)
        # new_node = tn.contract_between(nodes[a], nodes[b], allow_outer_product=True)
        # nodes = _multi_remove(nodes, [a, b])
        # nodes.insert(a, new_node)
    final_node = nodes[-1]
    if output_edge_order is not None:
        final_node.reorder_edges(output_edge_order)
    return final_node


# base = tn.contractors.opt_einsum_paths.path_contractors.base
# utils = tn.contractors.opt_einsum_paths.utils

_einsum_symbols_base = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"


@lru_cache(2**14)
def get_symbol(i: int) -> str:
    """Get the symbol corresponding to int ``i`` - runs through the usual 52
    letters before resorting to unicode characters, starting at ``chr(192)``
    and skipping surrogates. From cotengra codebase
    """
    if i < 52:
        # use a-z, A-Z first
        return _einsum_symbols_base[i]

    # then proceed from 'À'
    i += 140

    if i >= 55296:
        # Skip chr(57343) - chr(55296) as surrogates
        i += 2048

    return chr(i)


def _extract_topology(
    nodes: List[tn.Node],
) -> Tuple[List[Any], List[str], str, Dict[str, int]]:
    """
    Convert a physical tensor network graph (with possible CopyNodes) into
    algebraic components for einsum/cotengra.
    """
    # split nodes into regular nodes and CopyNodes
    nodes = sorted(nodes, key=lambda node: getattr(node, "_stable_id_", -1))
    regular_nodes = [n for n in nodes if not isinstance(n, tn.CopyNode)]
    copy_nodes = [n for n in nodes if isinstance(n, tn.CopyNode)]

    uf = UnionFind()
    all_edges = tn.get_all_edges(nodes)
    for edge in all_edges:
        uf[edge]  # init

    for cn in copy_nodes:
        edges = cn.edges
        if edges:
            root_edge = edges[0]
            for i in range(1, len(edges)):
                uf.union(root_edge, edges[i])

    mapping_dict = {}
    symbol_counter = 0

    input_sets = []
    raw_tensors = []
    for node in regular_nodes:
        node_symbols = []
        for edge in node.edges:
            root = uf[edge]
            if root not in mapping_dict:
                mapping_dict[root] = get_symbol(symbol_counter)
                symbol_counter += 1
            node_symbols.append(mapping_dict[root])
        input_sets.append("".join(node_symbols))
        raw_tensors.append(node.tensor)

    dangling_edges = sorted_edges(tn.get_subgraph_dangling(nodes))
    output_set = []
    for edge in dangling_edges:
        root = uf[edge]
        if root not in mapping_dict:
            mapping_dict[root] = get_symbol(symbol_counter)
            symbol_counter += 1
        output_set.append(mapping_dict[root])

    size_dict = {symbol: root.dimension for root, symbol in mapping_dict.items()}
    return (
        raw_tensors,
        input_sets,
        "".join(output_set),
        size_dict,
    )


def _normalize_omeco_topology(
    input_sets: Sequence[Sequence[Any]],
    output_set: Sequence[Any],
    size_dict: Dict[Any, int],
) -> Tuple[List[List[int]], List[int], Dict[int, int]]:
    mapping: Dict[Any, int] = {}

    def relabel(symbol: Any) -> int:
        if symbol not in mapping:
            mapping[symbol] = len(mapping)
        return mapping[symbol]

    new_inputs = [[relabel(symbol) for symbol in term] for term in input_sets]
    new_output = [relabel(symbol) for symbol in output_set]
    new_sizes = {relabel(symbol): int(size) for symbol, size in size_dict.items()}
    return new_inputs, new_output, new_sizes


def _omeco_tree_to_path(tree: Dict[str, Any], ntensors: int) -> List[Tuple[int, int]]:
    leaf_indices: List[int] = []

    def collect(node: Dict[str, Any]) -> None:
        if "tensor_index" in node:
            leaf_indices.append(int(node["tensor_index"]))
            return
        if "tensorindex" in node:
            leaf_indices.append(int(node["tensorindex"]))
            return
        for child in node["args"]:
            collect(child)

    collect(tree)
    offset = 1 if sorted(leaf_indices) == list(range(1, ntensors + 1)) else 0
    queue = list(range(ntensors))
    path: List[Tuple[int, int]] = []
    next_index = ntensors

    def visit(node: Dict[str, Any]) -> int:
        nonlocal next_index
        if "tensor_index" in node:
            return int(node["tensor_index"]) - offset
        if "tensorindex" in node:
            return int(node["tensorindex"]) - offset
        if len(node["args"]) != 2:
            raise ValueError("omeco returned a non-binary contraction tree")

        left = visit(node["args"][0])
        right = visit(node["args"][1])
        left_index = queue.index(left)
        right_index = queue.index(right)
        if left_index > right_index:
            pair = (left_index, right_index)
        else:
            pair = (right_index, left_index)
        queue.pop(pair[0])
        queue.pop(pair[1])
        queue.append(next_index)
        path.append(pair)
        next_index += 1
        return next_index - 1

    visit(tree)
    return path


class OMEOptimizer:
    """
    Adapt an ``omeco`` optimizer for ``tc.set_contractor("custom", ...)``.

    ``omeco`` stays an optional dependency and is imported lazily only when this
    adapter is actually used.

    Example:

    >>> import omeco
    >>> opt = omeco.TreeSA(ntrials=4, niters=8, betas=[0.1, 1.0, 10.0])
    >>> tc.set_contractor("custom", optimizer=opt, preprocessing=True)
    """

    def __init__(self, optimizer: Optional[Any] = None):
        self.optimizer = optimizer

    def __call__(
        self,
        input_sets: Sequence[Sequence[Any]],
        output_set: Sequence[Any],
        size_dict: Dict[Any, int],
        _memory_limit: Optional[int] = None,
        **_kws: Any,
    ) -> List[Tuple[int, int]]:
        try:
            import omeco
        except ImportError as e:
            raise ImportError(
                "OMEOptimizer requires the optional dependency 'omeco'."
            ) from e
        inputs, output, sizes = _normalize_omeco_topology(
            input_sets, output_set, size_dict
        )
        if self.optimizer is None:
            tree = omeco.optimize_code(inputs, output, sizes)
        else:
            tree = omeco.optimize_code(inputs, output, sizes, self.optimizer)
        return _omeco_tree_to_path(tree.to_dict(), len(inputs))


def _wrap_omeco_optimizer(optimizer: Any) -> Any:
    if optimizer is None or isinstance(optimizer, list) or callable(optimizer):
        return optimizer
    try:
        import omeco
    except ImportError:
        return optimizer
    if isinstance(optimizer, (omeco.GreedyMethod, omeco.TreeSA)):
        return OMEOptimizer(optimizer)
    return optimizer


def _algebraic_base_contraction(
    nodes: List[tn.Node],
    algorithm: Any,
    output_edge_order: Optional[Sequence[tn.Edge]] = None,
    ignore_edge_order: bool = False,
    **kws: Any,
) -> Any:
    """
    Execute contraction using cotengra and autoray for bare tensors.
    """
    import cotengra as ctg

    raw_tensors, input_sets, output_set, size_dict = _extract_topology(nodes)
    # Use the backend of the first node
    be = nodes[0].backend

    if len(raw_tensors) == 1:
        # Avoid cotengra bug for empty contraction paths
        final_raw_tensor = be.einsum(input_sets[0] + "->" + output_set, *raw_tensors)
        exponent = 0.0
    else:
        path = algorithm(input_sets, output_set, size_dict)
        logger.info("the contraction path is given as %s" % str(path))

        tree = ctg.ContractionTree.from_path(
            input_sets, output_set, size_dict, path=path
        )

        # Use autoray to keep AD and JIT support across backends
        # Note: cotengra's make_contractor handles the orchestration
        if not kws.get("strip_exponent", False):
            contractor = ctg.core.make_contractor(tree, implementation="autoray")
            final_raw_tensor = contractor(*raw_tensors)
        else:
            final_raw_tensor, exponent = tree.contract(raw_tensors, strip_exponent=True)

    final_node = tn.Node(final_raw_tensor, backend=be)

    # Resolve dangling edges in the same order as in _extract_topology
    dangling_edges = sorted_edges(tn.get_subgraph_dangling(nodes))

    # Update the edges to point to the new final_node
    for i, edge in enumerate(dangling_edges):
        if edge.node1 in nodes:
            edge.node1 = final_node
            edge.axis1 = i
        else:
            edge.node2 = final_node
            edge.axis2 = i
    final_node.edges = list(dangling_edges)

    if not ignore_edge_order:
        if output_edge_order is None:
            output_edge_order = dangling_edges
        final_node.reorder_edges(list(output_edge_order))

    if kws.get("strip_exponent", False):
        return final_node, exponent

    return final_node


def _get_path(
    nodes: List[tn.Node], algorithm: Any
) -> Tuple[List[Tuple[int, int]], List[tn.Node]]:
    nodes = list(nodes)
    input_sets = [set([id(e) for e in node.edges]) for node in nodes]
    output_set = set([id(e) for e in tn.get_subgraph_dangling(nodes)])
    size_dict = {id(edge): edge.dimension for edge in tn.get_all_edges(nodes)}

    return algorithm(input_sets, output_set, size_dict), nodes


def _identity(*args: Any, **kws: Any) -> Any:
    return args


def _get_path_cache_friendly(
    nodes: List[tn.Node], algorithm: Any
) -> Tuple[List[Tuple[int, int]], List[tn.Node]]:
    nodes = list(nodes)

    nodes_new = sorted(nodes, key=lambda node: getattr(node, "_stable_id_", -1))
    # if isinstance(algorithm, list):
    #     return algorithm, [nodes_new]

    all_edges = tn.get_all_edges(nodes_new)
    all_edges_sorted = sorted_edges(all_edges)
    mapping_dict = {}
    i = 0
    for edge in all_edges_sorted:
        if id(edge) not in mapping_dict:
            mapping_dict[id(edge)] = get_symbol(i)
            i += 1

    input_sets = [list([mapping_dict[id(e)] for e in node.edges]) for node in nodes_new]
    output_set = list(
        [mapping_dict[id(e)] for e in sorted_edges(tn.get_subgraph_dangling(nodes_new))]
    )
    size_dict = {mapping_dict[id(edge)]: edge.dimension for edge in all_edges_sorted}
    logger.debug("input_sets: %s" % input_sets)
    logger.debug("output_set: %s" % output_set)
    logger.debug("size_dict: %s" % size_dict)
    logger.debug("path finder algorithm: %s" % algorithm)
    return algorithm(input_sets, output_set, size_dict), nodes_new
    # directly get input_sets, output_set and size_dict by using identity function as algorithm


get_tn_info = partial(_get_path_cache_friendly, algorithm=_identity)


# some contractor setup usages
"""
import cotengra as ctg
import opt_einsum as oem

sys.setrecursionlimit(10000) # for successfullt ctg parallel

opt = ctg.ReusableHyperOptimizer(
    methods=["greedy", "kahypar"],
    parallel=True,
    minimize="write",
    max_time=30,
    max_repeats=4096,
    progbar=True,
)
tc.set_contractor("custom", optimizer=opt, preprocessing=True)
tc.set_contractor("custom_stateful", optimizer=oem.RandomGreedy, max_time=60, max_repeats=128, minimize="size")
tc.set_contractor("plain-experimental", local_steps=3)

# hyper efficient contractor: though long computation time required, suitable for extra large circuit simulation
opt = ctg.ReusableHyperOptimizer(
    minimize='combo',
    max_repeats=1024,
    max_time='equil:128',
    optlib='nevergrad',
    progbar=True,
)

def opt_reconf(inputs, output, size, **kws):
    tree = opt.search(inputs, output, size)
    tree_r = tree.subtree_reconfigure_forest(progbar=True, num_trees=10, 
                        num_restarts=20, subtree_weight_what=("size", ))
    return tree_r.get_path()

tc.set_contractor("custom", optimizer=opt_reconf)
"""


def _base(
    nodes: List[tn.Node],
    algorithm: Any,
    output_edge_order: Optional[Sequence[tn.Edge]] = None,
    ignore_edge_order: bool = False,
    total_size: Optional[int] = None,
    debug_level: int = 0,
    use_primitives: Optional[bool] = None,
    **kws: Any,
) -> Any:
    """
    The base method for all `opt_einsum` contractors.

    :param nodes: A collection of connected nodes.
    :type nodes: List[tn.Node]
    :pram algorithm: `opt_einsum` contraction method to use.
    :type algorithm: Any
    :param output_edge_order: An optional list of edges. Edges of the
        final node in `nodes_set` are reordered into `output_edge_order`;
        if final node has more than one edge, `output_edge_order` must be provided.
    :type output_edge_order: Optional[Sequence[tn.Edge]], optional
    :param ignore_edge_order: An option to ignore the output edge order.
    :type ignore_edge_order: bool
    :param total_size: The total size of the tensor network.
    :type total_size: Optional[int], optional
    :raises ValueError:"The final node after contraction has more than
        one remaining edge. In this case `output_edge_order` has to be provided," or
        "Output edges are not equal to the remaining non-contracted edges of the final node."
    :return: The final node after full contraction.
    :rtype: tn.Node
    """
    # rewrite tensornetwork default to add logging infras
    nodes_set = set(nodes)
    edges = tn.get_all_edges(nodes_set)
    # output edge order has to be determinded before any contraction
    # (edges are refreshed after contractions)

    if not ignore_edge_order:
        if output_edge_order is None:
            output_edge_order = list(tn.get_subgraph_dangling(nodes))
            if len(output_edge_order) > 1:
                raise ValueError(
                    "The final node after contraction has more than "
                    "one remaining edge. In this case `output_edge_order` "
                    "has to be provided."
                )

        if set(output_edge_order) != tn.get_subgraph_dangling(nodes):
            raise ValueError(
                "output edges are not equal to the remaining "
                "non-contracted edges of the final node."
            )

    # 1. Resolve topology and check for hyperedges
    has_hyperedges = any(isinstance(n, tn.CopyNode) for n in nodes)

    if use_primitives is True or (use_primitives is None and has_hyperedges):
        # ==========================================
        # NEW ALGEBRAIC EXECUTION PATH (Opt-in)
        # ==========================================
        return _algebraic_base_contraction(
            nodes, algorithm, output_edge_order, ignore_edge_order, **kws
        )

    # ==========================================
    # ORIGINAL EXECUTION PATH (100% Backward Compatible)
    # ==========================================

    for edge in edges:
        if not edge.is_disabled:  # if its disabled we already contracted it
            if edge.is_trace():
                idx = [i for i, n in enumerate(nodes) if id(n) == id(edge.node1)]
                nodes = _multi_remove(nodes, idx)
                nodes.append(tn.contract_parallel(edge))
                # nodes_set.remove(edge.node1)
                # nodes_set.add(tn.contract_parallel(edge))

    if len(nodes) == 1:
        # There's nothing to contract.
        if ignore_edge_order:
            return list(nodes)[0]
        return list(nodes)[0].reorder_edges(output_edge_order)

    # nodes = list(nodes_set)

    # Then apply `opt_einsum`'s algorithm
    # if isinstance(algorithm, list):
    #     path = algorithm
    # else:
    path, nodes = _get_path_cache_friendly(nodes, algorithm)
    if debug_level == 2:  # do nothing
        if output_edge_order:
            shape = [e.dimension for e in output_edge_order]
        else:
            shape = []
        return tn.Node(backend.zeros(shape))
    logger.info("the contraction path is given as %s" % str(path))
    if total_size is None:
        total_size = sum([_sizen(t) for t in nodes])
    for ab in path:
        if len(ab) < 2:
            logger.warning("single element tuple in contraction path!")
            continue
        a, b = ab

        if debug_level == 1:
            from .simplify import pseudo_contract_between

            new_node = pseudo_contract_between(nodes[a], nodes[b])
        else:
            new_node = tn.contract_between(nodes[a], nodes[b], allow_outer_product=True)
        nodes.append(new_node)
        # nodes[a] = backend.zeros([1])
        # nodes[b] = backend.zeros([1])
        nodes = _multi_remove(nodes, [a, b])

        logger.debug(_sizen(new_node, is_log=True))
        total_size += _sizen(new_node)
    logger.info("----- WRITE: %s --------\n" % np.log2(total_size))

    # if the final node has more than one edge,
    # output_edge_order has to be specified
    final_node = nodes[0]  # nodes were connected, we checked this
    if not ignore_edge_order:
        final_node.reorder_edges(output_edge_order)
    return final_node


class NodesReturn(Exception):
    """
    Intentionally stop execution to return a value.
    """

    def __init__(self, value_to_return: Any):
        self.value = value_to_return
        super().__init__(
            f"Intentionally stopping execution to return: {value_to_return}"
        )


def _get_sorted_nodes(nodes: List[Any], *args: Any, **kws: Any) -> Any:
    nodes_new = sorted(nodes, key=lambda node: getattr(node, "_stable_id_", -1))
    raise NodesReturn(nodes_new)


def function_nodes_capture(func: Callable[[Any], Any]) -> Callable[[Any], Any]:
    @wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        with runtime_contractor(method="before"):
            try:
                result = func(*args, **kwargs)
                return result
            except NodesReturn as e:
                return e.value

    return wrapper


@contextmanager
def runtime_nodes_capture(key: str = "nodes") -> Iterator[Any]:
    old_contractor = getattr(thismodule, "contractor")
    set_contractor(method="before")
    captured_value: Dict[str, List[tn.Node]] = {}
    try:
        yield captured_value
    except NodesReturn as e:
        captured_value[key] = e.value
    finally:
        _set_global_contractor(old_contractor)


def custom(
    nodes: List[Any],
    optimizer: Any,
    memory_limit: Optional[int] = None,
    output_edge_order: Optional[List[Any]] = None,
    ignore_edge_order: bool = False,
    debug_level: int = 0,
    use_primitives: Optional[bool] = None,
    **kws: Any,
) -> Any:
    local_kws = dict(kws)
    debug_level = local_kws.pop("debug_level", debug_level)
    if len(nodes) < 5:
        alg = opt_einsum.paths.optimal
        # not good at minimize WRITE actually...
        return _base(
            nodes,
            alg,
            output_edge_order,
            ignore_edge_order,
            debug_level=debug_level,
            use_primitives=use_primitives,
            **local_kws,
        )

    total_size = None
    has_hyperedges = any(isinstance(n, tn.CopyNode) for n in nodes)
    if local_kws.get("preprocessing", None) and not has_hyperedges:
        # nodes = _full_light_cone_cancel(nodes)
        nodes, total_size = _merge_single_gates(nodes)
    if not isinstance(optimizer, list):
        alg = partial(optimizer, memory_limit=memory_limit)
    else:
        alg = optimizer
    return _base(
        nodes,
        alg,
        output_edge_order,
        ignore_edge_order,
        total_size,
        debug_level=debug_level,
        use_primitives=use_primitives,
        **local_kws,
    )


def custom_stateful(
    nodes: List[Any],
    optimizer: Any,
    memory_limit: Optional[int] = None,
    opt_conf: Optional[Dict[str, Any]] = None,
    output_edge_order: Optional[List[Any]] = None,
    ignore_edge_order: bool = False,
    use_primitives: Optional[bool] = None,
    **kws: Any,
) -> Any:
    if opt_conf is None:
        opt_conf = {}
    opt = optimizer(**opt_conf)  # reinitiate the optimizer each time
    if kws.get("contraction_info", None):
        opt = contraction_info_decorator(opt)
    local_kws = dict(kws)
    debug_level = local_kws.pop("debug_level", 0)
    local_kws.pop("contraction_info", None)
    return custom(
        nodes,
        opt,
        memory_limit=memory_limit,
        output_edge_order=output_edge_order,
        ignore_edge_order=ignore_edge_order,
        debug_level=debug_level,
        use_primitives=use_primitives,
        **local_kws,
    )


# used by custom contractor variants
def contraction_info_decorator(algorithm: Callable[..., Any]) -> Callable[..., Any]:
    """Decorator to add contraction information logging to an optimizer.

    This decorator wraps an optimization algorithm and prints detailed information
    about the contraction cost (FLOPs, size, write) and path finding time.

    :param algorithm: The optimization algorithm to decorate.
    :type algorithm: Callable[..., Any]
    :return: The decorated optimization algorithm.
    :rtype: Callable[..., Any]
    """
    from cotengra import ContractionTree

    def new_algorithm(
        input_sets: Dict[Any, int],
        output_set: Dict[Any, int],
        size_dict: Dict[Any, int],
        **kws: Any,
    ) -> Any:
        t0 = time.time()
        path = algorithm(input_sets, output_set, size_dict, **kws)
        path_finding_time = time.time() - t0
        tree = ContractionTree.from_path(input_sets, output_set, size_dict, path=path)
        print("------ contraction cost summary ------")
        print(
            "log10[FLOPs]: ",
            "%.3f" % np.log10(float(tree.total_flops())),
            " log2[SIZE]: ",
            "%.0f" % tree.contraction_width(),
            " log2[WRITE]: ",
            "%.3f" % np.log2(float(tree.total_write())),
            " PathFindingTime: ",
            "%.3f" % path_finding_time,
        )
        return path

    return new_algorithm


def set_contractor(
    method: Optional[str] = None,
    optimizer: Optional[Any] = None,
    memory_limit: Optional[int] = None,
    opt_conf: Optional[Dict[str, Any]] = None,
    set_global: bool = True,
    contraction_info: bool = False,
    debug_level: int = 0,
    use_primitives: Optional[bool] = None,
    **kws: Any,
) -> Callable[..., Any]:
    """
    To set runtime contractor of the tensornetwork for a better contraction path.
    For more information on the usage of contractor, please refer to independent tutorial.

    :param method: "auto", "greedy", "branch", "plain", "tng", "custom", "custom_stateful". defaults to None ("auto")
    :type method: Optional[str], optional
    :param optimizer: Valid for "custom" or "custom_stateful" as method, defaults to None
    :type optimizer: Optional[Any], optional
    :param memory_limit: It is not very useful, as ``memory_limit`` leads to ``branch`` contraction
        instead of ``greedy`` which is rather slow, defaults to None
    :type memory_limit: Optional[int], optional
    :raises Exception: Tensornetwork version is too low to support some of the contractors.
    :raises ValueError: Unknown method options.
    :return: The new tensornetwork with its contractor set.
    :rtype: tn.Node
    """
    if not method:
        method = "greedy"
        # auto for small size fallbacks to dp, which has bug for now
        # see: https://github.com/dgasmith/opt_einsum/issues/172
    if kws.get("strip_exponent", False):
        if use_primitives is None:
            use_primitives = True
    if method.startswith("cotengra"):
        # cotengra shortcut
        import cotengra

        if method == "cotengra":
            method = "custom"
            optimizer = cotengra.ReusableHyperOptimizer(
                methods=["greedy", "kahypar"],
                parallel=True,
                minimize="combo",
                max_time=30,
                max_repeats=64,
                progbar=True,
            )
        else:  # "cotengra-30-64"
            _, mt, mr = method.split("-")
            method = "custom"
            optimizer = cotengra.ReusableHyperOptimizer(
                methods=["greedy", "kahypar"],
                parallel=True,
                minimize="combo",
                max_time=int(mt),
                max_repeats=int(mr),
                progbar=True,
            )

    if method == "plain":
        cf = plain_contractor
    elif method == "plain-experimental":
        cf = partial(experimental_contractor, local_steps=kws.get("local_steps", 2))
    elif method == "tng":  # don't use, deprecated, no guarantee
        if has_ps:
            cf = tn_greedy_contractor
        else:
            raise Exception(
                "current version of tensornetwork doesn't support tng contraction"
            )
    elif method == "custom_stateful":
        cf = custom_stateful  # type: ignore
        cf = partial(  # type: ignore
            cf,
            optimizer=optimizer,
            opt_conf=opt_conf,
            contraction_info=contraction_info,
            debug_level=debug_level,
            use_primitives=use_primitives,
            **kws,
        )

    elif method == "before":  # a hack way to get the nodes
        cf = _get_sorted_nodes

    else:
        # cf = getattr(tn.contractors, method, None)
        # if not cf:
        # raise ValueError("Unknown contractor type: %s" % method)

        if method != "custom":
            optimizer = getattr(opt_einsum.paths, method)
        else:
            optimizer = _wrap_omeco_optimizer(optimizer)
        if contraction_info is True:
            optimizer = contraction_info_decorator(optimizer)  # type: ignore
        cf = partial(
            custom,
            optimizer=optimizer,
            memory_limit=memory_limit,
            debug_level=debug_level,
            use_primitives=use_primitives,
            **kws,
        )
    if set_global:
        _set_global_contractor(cf)
    return cf


set_contractor("greedy", preprocessing=True)

get_contractor = partial(set_contractor, set_global=False)


def set_function_contractor(*confargs: Any, **confkws: Any) -> Callable[..., Any]:
    """
    Function decorate to change function-level contractor

    :return: _description_
    :rtype: Callable[..., Any]
    """

    def wrapper(f: Callable[..., Any]) -> Callable[..., Any]:
        @wraps(f)
        def newf(*args: Any, **kws: Any) -> Any:
            old_contractor = getattr(thismodule, "contractor")
            set_contractor(*confargs, **confkws)
            try:
                return f(*args, **kws)
            finally:
                _set_global_contractor(old_contractor)

        return newf

    return wrapper


@contextmanager
def runtime_contractor(*confargs: Any, **confkws: Any) -> Iterator[Any]:
    """
    Context manager to change with-levek contractor

    :yield: _description_
    :rtype: Iterator[Any]
    """
    old_contractor = getattr(thismodule, "contractor")
    nc = set_contractor(*confargs, **confkws)
    try:
        yield nc
    finally:
        _set_global_contractor(old_contractor)


def split_rules(
    max_singular_values: Optional[int] = None,
    max_truncation_err: Optional[float] = None,
    relative: bool = False,
) -> Any:
    """
    Obtain the direcionary of truncation rules

    :param max_singular_values: The maximum number of singular values to keep.
    :type max_singular_values: int, optional
    :param max_truncation_err: The maximum allowed truncation error.
    :type max_truncation_err: float, optional
    :param relative: Multiply `max_truncation_err` with the largest singular value.
    :type relative: bool, optional
    """
    rules: Any = {}
    if max_singular_values is not None:
        rules["max_singular_values"] = max_singular_values
    if max_truncation_err is not None:
        rules["max_truncattion_err"] = max_truncation_err
    if relative is not None:
        rules["relative"] = relative
    return rules
