# Copyright 2019 The TensorNetwork Authors
#
# 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.
"""Implementation of TensorNetwork structure."""

import collections
from typing import Any, Dict, List, Optional, Set, Text, Tuple, Union, \
    Sequence, Iterable, Type
import numpy as np

#pylint: disable=useless-import-alias
#pylint: disable=line-too-long
from tensornetwork.network_components import BaseNode, Node, CopyNode, Edge, disconnect
from tensornetwork.backends import backend_factory
from tensornetwork.backends.base_backend import BaseBackend
from tensornetwork.network_components import connect, contract_parallel
Tensor = Any


def norm(node: BaseNode) -> Tensor:
  """The L2 norm of `node`
  Args:
    node: A `BaseNode`. 
  Returns:
    The L2 norm 
  Raises:
    AttributeError: If `node` has no `backend` attribute.
  """
  if not hasattr(node, 'backend'):
    raise AttributeError('Node {} of type {} has no `backend`'.format(
        node, type(node)))
  return node.backend.norm(node.tensor)


def conj(node: BaseNode,
         name: Optional[Text] = None,
         axis_names: Optional[List[Text]] = None) -> BaseNode:
  """Conjugate `node`.

  Args:
    node: A `BaseNode`. 
    name: Optional name to give the new node.
    axis_names: Optional list of names for the axis.
  Returns:
    A new node. The complex conjugate of `node`.
  Raises:
    AttributeError: If `node` has no `backend` attribute.
  """
  if not hasattr(node, 'backend'):
    raise AttributeError('Node {} of type {} has no `backend`'.format(
        node, type(node)))
  backend = node.backend
  if not axis_names:
    axis_names = node.axis_names

  return Node(
      backend.conj(node.tensor),
      name=name,
      axis_names=axis_names,
      backend=backend)


def transpose(node: BaseNode,
              permutation: Sequence[Union[Text, int]],
              name: Optional[Text] = None,
              axis_names: Optional[List[Text]] = None) -> BaseNode:
  """Transpose `node`
  Args:
    node: A `BaseNode`. 
    permutation: A list of int or str. The permutation of the axis
    name: Optional name to give the new node.
    axis_names: Optional list of names for the axis.
  Returns:
    A new node. The transpose of `node`.
  Raises:
    AttributeError: If `node` has no `backend` attribute, or if 
      `node` has no tensor.
    ValueError: If either `permutation` is not the same as expected or
      if you try to permute with a trace edge.
  """

  if not hasattr(node, 'backend'):
    raise AttributeError('Node {} of type {} has no `backend`'.format(
        node, type(node)))

  perm = [node.get_axis_number(p) for p in permutation]
  if not axis_names:
    axis_names = node.axis_names

  new_node = Node(
      node.tensor,
      name=name,
      axis_names=node.axis_names,
      backend=node.backend)
  return new_node.reorder_axes(perm)


def copy(nodes: Iterable[BaseNode],
         conjugate: bool = False) -> Tuple[dict, dict]:
  """Copy the given nodes and their edges.

  This will return a dictionary linking original nodes/edges 
  to their copies. If nodes A and B are connected but only A is passed in to be
  copied, the edge between them will become a dangling edge.

  Args:
    nodes: An `Iterable` (Usually a `List` or `Set`) of `Nodes`.
    conjugate: Boolean. Whether to conjugate all of the nodes
        (useful for calculating norms and reduced density
        matrices).

  Returns:
    A tuple containing:
      node_dict: A dictionary mapping the nodes to their copies.
      edge_dict: A dictionary mapping the edges to their copies.
  """
  node_dict = {}
  for node in nodes:
    if isinstance(node, CopyNode):
      node_dict[node] = CopyNode(
          node.rank,
          node.dimension,
          name=node.name,
          axis_names=node.axis_names,
          backend=node.backend,
          dtype=node.dtype)
    else:
      if conjugate:
        node_dict[node] = Node(
            node.backend.conj(node.tensor),
            name=node.name,
            axis_names=node.axis_names,
            backend=node.backend)
      else:
        node_dict[node] = Node(
            node.tensor,
            name=node.name,
            axis_names=node.axis_names,
            backend=node.backend)
  edge_dict = {}
  for edge in get_all_edges(nodes):
    node1 = edge.node1
    axis1 = edge.node1.get_axis_number(edge.axis1)
    # edge dangling or node2 does not need to be copied
    if edge.is_dangling() or edge.node2 not in node_dict:
      new_edge = Edge(node_dict[node1], axis1, edge.name)
      node_dict[node1].add_edge(new_edge, axis1)
      edge_dict[edge] = new_edge
      continue

    node2 = edge.node2
    axis2 = edge.node2.get_axis_number(edge.axis2)
    # copy node2 but not node1
    if node1 not in node_dict:
      new_edge = Edge(node_dict[node2], axis2, edge.name)
      node_dict[node2].add_edge(new_edge, axis2)
      edge_dict[edge] = new_edge
      continue

    # both nodes should be copied
    new_edge = Edge(node_dict[node1], axis1, edge.name, node_dict[node2], axis2)
    new_edge.set_signature(edge.signature)
    node_dict[node2].add_edge(new_edge, axis2)
    node_dict[node1].add_edge(new_edge, axis1)
    edge_dict[edge] = new_edge

  return node_dict, edge_dict


def remove_node(node: BaseNode) -> Tuple[Dict[Text, Edge], Dict[int, Edge]]:
  """Remove a node from the network.

  Args:
    node: The node to be removed.

  Returns:
    A tuple of:
      disconnected_edges_by_name: A Dictionary mapping `node`'s axis names to
        the newly broken edges.
      disconnected_edges_by_axis: A Dictionary mapping `node`'s axis numbers
        to the newly broken edges.
  """
  disconnected_edges_by_name = {}
  disconnected_edges_by_axis = {}
  for i, name in enumerate(node.axis_names):
    if not node[i].is_dangling() and not node[i].is_trace():
      edge1, edge2 = disconnect(node[i])
      new_disconnected_edge = edge1 if edge1.node1 is not node else edge2
      disconnected_edges_by_axis[i] = new_disconnected_edge
      disconnected_edges_by_name[name] = new_disconnected_edge
  return disconnected_edges_by_name, disconnected_edges_by_axis


def split_node(
    node: BaseNode,
    left_edges: List[Edge],
    right_edges: List[Edge],
    max_singular_values: Optional[int] = None,
    max_truncation_err: Optional[float] = None,
    left_name: Optional[Text] = None,
    right_name: Optional[Text] = None,
    edge_name: Optional[Text] = None,
) -> Tuple[BaseNode, BaseNode, Tensor]:
  """Split a `Node` using Singular Value Decomposition.

  Let M be the matrix created by flattening left_edges and right_edges into
  2 axes. Let :math:`U S V^* = M` be the Singular Value Decomposition of 
  :math:`M`. This will split the network into 2 nodes. The left node's 
  tensor will be :math:`U \\sqrt{S}` and the right node's tensor will be 
  :math:`\\sqrt{S} V^*` where :math:`V^*` is
  the adjoint of :math:`V`.

  The singular value decomposition is truncated if `max_singular_values` or
  `max_truncation_err` is not `None`.

  The truncation error is the 2-norm of the vector of truncated singular
  values. If only `max_truncation_err` is set, as many singular values will
  be truncated as possible while maintaining:
  `norm(truncated_singular_values) <= max_truncation_err`.

  If only `max_singular_values` is set, the number of singular values kept
  will be `min(max_singular_values, number_of_singular_values)`, so that
  `max(0, number_of_singular_values - max_singular_values)` are truncated.

  If both `max_truncation_err` and `max_singular_values` are set,
  `max_singular_values` takes priority: The truncation error may be larger
  than `max_truncation_err` if required to satisfy `max_singular_values`.

  Args:
    node: The node you want to split.
    left_edges: The edges you want connected to the new left node.
    right_edges: The edges you want connected to the new right node.
    max_singular_values: The maximum number of singular values to keep.
    max_truncation_err: The maximum allowed truncation error.
    left_name: The name of the new left node. If `None`, a name will be 
      generated automatically.
    right_name: The name of the new right node. If `None`, a name will be 
      generated automatically.
    edge_name: The name of the new `Edge` connecting the new left and 
      right node. If `None`, a name will be generated automatically. 
      The new axis will get the same name as the edge.

  Returns:
    A tuple containing:
      left_node: 
        A new node created that connects to all of the `left_edges`.
        Its underlying tensor is :math:`U \\sqrt{S}`
      right_node: 
        A new node created that connects to all of the `right_edges`.
        Its underlying tensor is :math:`\\sqrt{S} V^*`
      truncated_singular_values: 
        The vector of truncated singular values.
  Raises:
    AttributeError: If `node` has no backend attribute
  """

  if not hasattr(node, 'backend'):
    raise AttributeError('Node {} of type {} has no `backend`'.format(
        node, type(node)))

  if node.axis_names and edge_name:
    left_axis_names = []
    right_axis_names = [edge_name]
    for edge in left_edges:
      left_axis_names.append(node.axis_names[edge.axis1] if edge.node1 is node
                             else node.axis_names[edge.axis2])
    for edge in right_edges:
      right_axis_names.append(node.axis_names[edge.axis1] if edge.node1 is node
                              else node.axis_names[edge.axis2])
    left_axis_names.append(edge_name)
  else:
    left_axis_names = None
    right_axis_names = None

  backend = node.backend
  node.reorder_edges(left_edges + right_edges)

  u, s, vh, trun_vals = backend.svd_decomposition(node.tensor, len(left_edges),
                                                  max_singular_values,
                                                  max_truncation_err)
  sqrt_s = backend.sqrt(s)
  u_s = u * sqrt_s
  # We have to do this since we are doing element-wise multiplication against
  # the first axis of vh. If we don't, it's possible one of the other axes of
  # vh will be the same size as sqrt_s and would multiply across that axis
  # instead, which is bad.
  sqrt_s_broadcast_shape = backend.shape_concat(
      [backend.shape_tensor(sqrt_s), [1] * (len(vh.shape) - 1)], axis=-1)
  vh_s = vh * backend.reshape(sqrt_s, sqrt_s_broadcast_shape)
  left_node = Node(
      u_s, name=left_name, axis_names=left_axis_names, backend=backend)
  for i, edge in enumerate(left_edges):
    left_node.add_edge(edge, i)
    edge.update_axis(i, node, i, left_node)
  right_node = Node(
      vh_s, name=right_name, axis_names=right_axis_names, backend=backend)
  for i, edge in enumerate(right_edges):
    # i + 1 to account for the new edge.
    right_node.add_edge(edge, i + 1)
    edge.update_axis(i + len(left_edges), node, i + 1, right_node)
  connect(left_node.edges[-1], right_node.edges[0], name=edge_name)
  node.fresh_edges(node.axis_names)
  return left_node, right_node, trun_vals


def split_node_qr(
    node: BaseNode,
    left_edges: List[Edge],
    right_edges: List[Edge],
    left_name: Optional[Text] = None,
    right_name: Optional[Text] = None,
    edge_name: Optional[Text] = None,
) -> Tuple[BaseNode, BaseNode]:
  """Split a `Node` using QR decomposition

  Let M be the matrix created by flattening left_edges and right_edges into
  2 axes. Let :math:`QR = M` be the QR Decomposition of
  :math:`M`. This will split the network into 2 nodes. The left node's
  tensor will be :math:`Q` (an orthonormal matrix) and the right node's 
  tensor will be :math:`R` (an upper triangular matrix)

  Args:
    node: The node you want to split.
    left_edges: The edges you want connected to the new left node.
    right_edges: The edges you want connected to the new right node.
    left_name: The name of the new left node. If `None`, a name will be 
      generated automatically.
    right_name: The name of the new right node. If `None`, a name will be 
      generated automatically.
    edge_name: The name of the new `Edge` connecting the new left and right 
      node. If `None`, a name will be generated automatically.

  Returns:
    A tuple containing:
      left_node:
        A new node created that connects to all of the `left_edges`.
        Its underlying tensor is :math:`Q`
      right_node:
        A new node created that connects to all of the `right_edges`.
        Its underlying tensor is :math:`R`
  Raises:
    AttributeError: If `node` has no backend attribute
  """
  if not hasattr(node, 'backend'):
    raise AttributeError('Node {} of type {} has no `backend`'.format(
        node, type(node)))

  if node.axis_names and edge_name:
    left_axis_names = []
    right_axis_names = [edge_name]
    for edge in left_edges:
      left_axis_names.append(node.axis_names[edge.axis1] if edge.node1 is node
                             else node.axis_names[edge.axis2])
    for edge in right_edges:
      right_axis_names.append(node.axis_names[edge.axis1] if edge.node1 is node
                              else node.axis_names[edge.axis2])
    left_axis_names.append(edge_name)
  else:
    left_axis_names = None
    right_axis_names = None

  backend = node.backend
  node.reorder_edges(left_edges + right_edges)
  q, r = backend.qr_decomposition(node.tensor, len(left_edges))
  left_node = Node(
      q, name=left_name, axis_names=left_axis_names, backend=backend)
  for i, edge in enumerate(left_edges):
    left_node.add_edge(edge, i)
    edge.update_axis(i, node, i, left_node)
  right_node = Node(
      r, name=right_name, axis_names=right_axis_names, backend=backend)
  for i, edge in enumerate(right_edges):
    # i + 1 to account for the new edge.
    right_node.add_edge(edge, i + 1)
    edge.update_axis(i + len(left_edges), node, i + 1, right_node)
  connect(left_node.edges[-1], right_node.edges[0], name=edge_name)
  return left_node, right_node


def split_node_rq(
    node: BaseNode,
    left_edges: List[Edge],
    right_edges: List[Edge],
    left_name: Optional[Text] = None,
    right_name: Optional[Text] = None,
    edge_name: Optional[Text] = None,
) -> Tuple[BaseNode, BaseNode]:
  """Split a `Node` using RQ (reversed QR) decomposition

  Let M be the matrix created by flattening left_edges and right_edges into
  2 axes. Let :math:`QR = M^*` be the QR Decomposition of
  :math:`M^*`. This will split the network into 2 nodes. The left node's
  tensor will be :math:`R^*` (a lower triangular matrix) and the right 
  node's tensor will be :math:`Q^*` (an orthonormal matrix)

  Args:
    node: The node you want to split.
    left_edges: The edges you want connected to the new left node.
    right_edges: The edges you want connected to the new right node.
    left_name: The name of the new left node. If `None`, a name will be 
      generated automatically.
    right_name: The name of the new right node. If `None`, a name will be 
      generated automatically.
    edge_name: The name of the new `Edge` connecting the new left and 
      right node. If `None`, a name will be generated automatically.

  Returns:
    A tuple containing:
      left_node:
        A new node that connects to all of the `left_edges`.
        Its underlying tensor is :math:`R^*`
      right_node:
        A new node that connects to all of the `right_edges`.
        Its underlying tensor is :math:`Q^*`
  Raises:
    AttributeError: If `node` has no backend attribute

  """
  if not hasattr(node, 'backend'):
    raise AttributeError('Node {} of type {} has no `backend`'.format(
        node, type(node)))

  if node.axis_names and edge_name:
    left_axis_names = []
    right_axis_names = [edge_name]
    for edge in left_edges:
      left_axis_names.append(node.axis_names[edge.axis1] if edge.node1 is node
                             else node.axis_names[edge.axis2])
    for edge in right_edges:
      right_axis_names.append(node.axis_names[edge.axis1] if edge.node1 is node
                              else node.axis_names[edge.axis2])
    left_axis_names.append(edge_name)
  else:
    left_axis_names = None
    right_axis_names = None
  backend = node.backend
  node.reorder_edges(left_edges + right_edges)
  r, q = backend.rq_decomposition(node.tensor, len(left_edges))
  left_node = Node(
      r, name=left_name, axis_names=left_axis_names, backend=backend)
  for i, edge in enumerate(left_edges):
    left_node.add_edge(edge, i)
    edge.update_axis(i, node, i, left_node)
  right_node = Node(
      q, name=right_name, axis_names=right_axis_names, backend=backend)
  for i, edge in enumerate(right_edges):
    # i + 1 to account for the new edge.
    right_node.add_edge(edge, i + 1)
    edge.update_axis(i + len(left_edges), node, i + 1, right_node)
  connect(left_node.edges[-1], right_node.edges[0], name=edge_name)
  return left_node, right_node


def split_node_full_svd(
    node: BaseNode,
    left_edges: List[Edge],
    right_edges: List[Edge],
    max_singular_values: Optional[int] = None,
    max_truncation_err: Optional[float] = None,
    left_name: Optional[Text] = None,
    middle_name: Optional[Text] = None,
    right_name: Optional[Text] = None,
    left_edge_name: Optional[Text] = None,
    right_edge_name: Optional[Text] = None,
) -> Tuple[BaseNode, BaseNode, BaseNode, Tensor]:
  """Split a node by doing a full singular value decomposition.

  Let M be the matrix created by flattening left_edges and right_edges into
  2 axes. Let :math:`U S V^* = M` be the Singular Value Decomposition of
  :math:`M`.

  The left most node will be :math:`U` tensor of the SVD, the middle node is
  the diagonal matrix of the singular values, ordered largest to smallest,
  and the right most node will be the :math:`V*` tensor of the SVD.

  The singular value decomposition is truncated if `max_singular_values` or
  `max_truncation_err` is not `None`.

  The truncation error is the 2-norm of the vector of truncated singular
  values. If only `max_truncation_err` is set, as many singular values will
  be truncated as possible while maintaining:
  `norm(truncated_singular_values) <= max_truncation_err`.

  If only `max_singular_values` is set, the number of singular values kept
  will be `min(max_singular_values, number_of_singular_values)`, so that
  `max(0, number_of_singular_values - max_singular_values)` are truncated.

  If both `max_truncation_err` and `max_singular_values` are set,
  `max_singular_values` takes priority: The truncation error may be larger
  than `max_truncation_err` if required to satisfy `max_singular_values`.

  Args:
    node: The node you want to split.
    left_edges: The edges you want connected to the new left node.
    right_edges: The edges you want connected to the new right node.
    max_singular_values: The maximum number of singular values to keep.
    max_truncation_err: The maximum allowed truncation error.
    left_name: The name of the new left node. If None, a name will be 
      generated automatically.
    middle_name: The name of the new center node. If None, a name will be 
      generated automatically.
    right_name: The name of the new right node. If None, a name will be 
      generated automatically.
    left_edge_name: The name of the new left `Edge` connecting
      the new left node (`U`) and the new central node (`S`).
      If `None`, a name will be generated automatically.
    right_edge_name: The name of the new right `Edge` connecting
      the new central node (`S`) and the new right node (`V*`).
      If `None`, a name will be generated automatically.

  Returns:
    A tuple containing:
      left_node:
        A new node created that connects to all of the `left_edges`.
        Its underlying tensor is :math:`U`
      singular_values_node:
        A new node that has 2 edges connecting `left_node` and `right_node`.
        Its underlying tensor is :math:`S`
      right_node:
        A new node created that connects to all of the `right_edges`.
        Its underlying tensor is :math:`V^*`
      truncated_singular_values:
        The vector of truncated singular values.
  Raises:
    AttributeError: If `node` has no backend attribute
  """
  if not hasattr(node, 'backend'):
    raise AttributeError('Node {} of type {} has no `backend`'.format(
        node, type(node)))

  if node.axis_names and left_edge_name and right_edge_name:
    left_axis_names = []
    right_axis_names = [right_edge_name]
    for edge in left_edges:
      left_axis_names.append(node.axis_names[edge.axis1] if edge.node1 is node
                             else node.axis_names[edge.axis2])
    for edge in right_edges:
      right_axis_names.append(node.axis_names[edge.axis1] if edge.node1 is node
                              else node.axis_names[edge.axis2])
    left_axis_names.append(left_edge_name)
    center_axis_names = [left_edge_name, right_edge_name]
  else:
    left_axis_names = None
    center_axis_names = None
    right_axis_names = None

  backend = node.backend

  node.reorder_edges(left_edges + right_edges)
  u, s, vh, trun_vals = backend.svd_decomposition(node.tensor, len(left_edges),
                                                  max_singular_values,
                                                  max_truncation_err)
  left_node = Node(
      u, name=left_name, axis_names=left_axis_names, backend=backend)
  singular_values_node = Node(
      backend.diag(s),
      name=middle_name,
      axis_names=center_axis_names,
      backend=backend)

  right_node = Node(
      vh, name=right_name, axis_names=right_axis_names, backend=backend)

  for i, edge in enumerate(left_edges):
    left_node.add_edge(edge, i)
    edge.update_axis(i, node, i, left_node)
  for i, edge in enumerate(right_edges):
    # i + 1 to account for the new edge.
    right_node.add_edge(edge, i + 1)
    edge.update_axis(i + len(left_edges), node, i + 1, right_node)
  connect(
      left_node.edges[-1], singular_values_node.edges[0], name=left_edge_name)
  connect(
      singular_values_node.edges[1], right_node.edges[0], name=right_edge_name)
  return left_node, singular_values_node, right_node, trun_vals


def _reachable(nodes: Set[BaseNode]) -> Set[BaseNode]:
  if not nodes:
    raise ValueError("Reachable requires at least 1 node.")
  node_que = collections.deque(nodes)
  seen_nodes = set()
  while node_que:
    node = node_que.popleft()
    if node not in seen_nodes:
      seen_nodes.add(node)
    for e in node.edges:
      for n in e.get_nodes():
        if n is not None and n not in seen_nodes:
          node_que.append(n)
          seen_nodes.add(n)
  return seen_nodes


def reachable(inputs: Union[BaseNode, Iterable[BaseNode], Edge, Iterable[Edge]]
             ) -> Set[BaseNode]:
  """
  Computes all nodes reachable from `node` or `edge.node1` by connected edges.
  Args:
    inputs: A `BaseNode`/`Edge` or collection of `BaseNodes`/`Edges`
  Returns:
    A set of `BaseNode` objects that can be reached from `node`
    via connected edges.
  Raises:
    ValueError: If an unknown value for `strategy` is passed.
  """

  if isinstance(inputs, BaseNode):
    inputs = {inputs}
  elif isinstance(inputs, Edge):
    inputs = {inputs.node1}
  elif isinstance(inputs, list) and all(isinstance(x, Edge) for x in inputs):
    inputs = {x.node1 for x in inputs}
  return _reachable(set(inputs))


def check_correct(nodes: Iterable[BaseNode],
                  check_connections: Optional[bool] = True) -> None:
  """
  Check if the network defined by `nodes` fulfills necessary
  consistency relations.

  Args:
    nodes: A list of `BaseNode` objects.
    check_connections: Check if the network is connected.
  Returns:
    None
  Raises:
    ValueError: If the network defined by `nodes` is not 
      correctly structured.
  """
  for node in nodes:
    for i, edge in enumerate(node.edges):
      if edge.node1 is not node and edge.node2 is not node:
        raise ValueError("Edge '{}' does not connect to node '{}'."
                         "Edge's nodes: '{}', '{}'.".format(
                             edge, node, edge.node1, edge.node2))

      is_edge_node_consistent = False
      if edge.node1 is node:
        if edge.axis1 == i:
          is_edge_node_consistent = True
      if edge.node2 is node:
        if edge.axis2 == i:
          is_edge_node_consistent = True
      if not is_edge_node_consistent:
        raise ValueError(
            "Edge '{}' does not point to '{}' on the correct axis. "
            "Edge axes: {}, {}. Node axis: {}.".format(edge, node, edge.axis1,
                                                       edge.axis2, i))
  if check_connections:
    check_connected(nodes)


def check_connected(nodes: Iterable[BaseNode]) -> None:
  """
  Check if all nodes in `nodes` are connected.
  Args:
    nodes: A list of nodes.
  Returns:
    None:
  Raises:
    ValueError: If not all nodes in `nodes` are connected.
  """
  nodes = list(nodes)
  if not set(nodes) <= reachable([nodes[0]]):
    raise ValueError("Non-connected graph")


def get_all_nodes(edges: Iterable[Edge]) -> Set[BaseNode]:
  """Return the set of nodes connected to edges."""
  nodes = set()
  for edge in edges:
    if edge.node1 is not None:
      nodes |= {edge.node1}
    if edge.node2 is not None:
      nodes |= {edge.node2}

  return nodes


def get_all_edges(nodes: Iterable[BaseNode]) -> Set[Edge]:
  """Return the set of edges of all nodes."""
  edges = set()
  for node in nodes:
    edges |= set(node.edges)
  return edges


def get_subgraph_dangling(nodes: Iterable[BaseNode]) -> Set[Edge]:
  """Get all of the edges that are "relatively dangling" to the given nodes.

  A "relatively dangling" edge is an edge that is either actually dangling
  or is connected to another node that is outside of the given collection
  of `nodes`.
  
  Args:
    nodes: A set of nodes.

  Returns:
    The set of "relatively dangling edges.
  """
  output = set()
  for edge in get_all_edges(nodes):
    if edge.is_dangling() or not set(edge.get_nodes()) <= set(nodes):
      output.add(edge)
  return output


def contract_trace_edges(node: BaseNode) -> BaseNode:
  """
  contract all trace edges of `node`.
  Args:
    node: A `BaseNode` object
  Returns:
    A new `BaseNode` obtained from contracting all 
    trace edges.
  Raises:
    ValueError: If `node` has no trace edges
  """
  for edge in node.edges:
    if edge.is_trace():
      return contract_parallel(edge)
  raise ValueError('`node` has no trace edges')


def reduced_density(traced_out_edges: Iterable[Edge]) -> Tuple[dict, dict]:
  """
  Constructs the tensor network for a reduced density matrix, given a pure state.

  The tensor network connected to `traced_out_edges` is assumed to be a pure
  quantum state (a state vector). This modifies the network so that it
  describes the reduced density matrix obtained by "tracing out" the specified
  edges.

  This is done by making a conjugate copy of the original network and
  connecting each edge in `traced_out_edges` with its conjugate counterpart.

  The edges in `edge_dict` corresponding to `traced_out_edges` will be the
  new non-dangling edges connecting the state with its conjugate.

  Args:
    traced_out_edges: A list of dangling edges
  Returns:
    A tuple containing:
      node_dict: A dictionary mapping the nodes in the original network to
        their conjugate copies.
      edge_dict: A dictionary mapping edges in the original network to their
        conjugate copies.
  """

  if list(filter(lambda x: not x.is_dangling(), traced_out_edges)):
    raise ValueError("traced_out_edges must only include dangling edges!")

  # Get all reachable nodes.
  old_nodes = reachable(get_all_nodes(traced_out_edges))

  # Copy and conjugate all reachable nodes.
  node_dict, edge_dict = copy(old_nodes, True)
  for t_edge in traced_out_edges:
    # Add each edge to the copied nodes as new edge.
    edge_dict[t_edge] = edge_dict[t_edge] ^ t_edge

  return node_dict, edge_dict


def switch_backend(nodes: Iterable[BaseNode], new_backend: Text) -> None:
  """Change the backend of the nodes.

  This will convert all node's tensors to the new backend's Tensor type.
  Args:
    nodes: iterable of nodes
    new_backend (str): The new backend.
    dtype (datatype): The dtype of the backend. If None, a defautl dtype according
                       to config.py will be chosen.
  """
  backend = backend_factory.get_backend(new_backend)
  for node in nodes:
    if node.backend.name != "numpy":
      raise NotImplementedError("Can only switch backends when the current "
                                "backend is 'numpy'. Current backend "
                                "is '{}'".format(node.backend))
    node.tensor = backend.convert_to_tensor(node.tensor)
    node.backend = backend
