# This code is part of Qiskit.
#
# (C) Copyright IBM 2023.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""This module contains common utils for disjoint coupling maps."""
from __future__ import annotations
from collections import defaultdict
from typing import List, Callable, TypeVar, Dict, Union
import uuid

import rustworkx as rx
from qiskit.dagcircuit import DAGOpNode

from qiskit.circuit import Qubit, Barrier, Clbit
from qiskit.dagcircuit.dagcircuit import DAGCircuit
from qiskit.dagcircuit.dagnode import DAGOutNode
from qiskit.transpiler.coupling import CouplingMap
from qiskit.transpiler.target import Target
from qiskit.transpiler.exceptions import TranspilerError
from qiskit.transpiler.passes.layout import vf2_utils

T = TypeVar("T")


def run_pass_over_connected_components(
    dag: DAGCircuit,
    components_source: Union[Target, CouplingMap],
    run_func: Callable[[DAGCircuit, CouplingMap], T],
) -> List[T]:
    """Run a transpiler pass inner function over mapped components."""
    if isinstance(components_source, Target):
        coupling_map = components_source.build_coupling_map(filter_idle_qubits=True)
    else:
        coupling_map = components_source
    cmap_components = coupling_map.connected_components()
    # If graph is connected we only need to run the pass once
    if len(cmap_components) == 1:
        if dag.num_qubits() > cmap_components[0].size():
            raise TranspilerError(
                "A connected component of the DAGCircuit is too large for any of the connected "
                "components in the coupling map."
            )
        return [run_func(dag, cmap_components[0])]
    dag_components = separate_dag(dag)
    mapped_components = map_components(dag_components, cmap_components)
    out_component_pairs = []
    for cmap_index, dags in mapped_components.items():
        # Take the first dag from the mapped dag components and then
        # compose it with any other dag components that are operating on the
        # same coupling map connected component. This results in a subcircuit
        # of possibly disjoint circuit components which we will run the layout
        # pass on.
        out_dag = dag_components[dags.pop()]
        for dag_index in dags:
            dag = dag_components[dag_index]
            out_dag.add_qubits(dag.qubits)
            out_dag.add_clbits(dag.clbits)
            for qreg in dag.qregs:
                out_dag.add_qreg(qreg)
            for creg in dag.cregs:
                out_dag.add_creg(creg)
            out_dag.compose(dag, qubits=dag.qubits, clbits=dag.clbits)
        out_component_pairs.append((out_dag, cmap_components[cmap_index]))
    res = [run_func(out_dag, cmap) for out_dag, cmap in out_component_pairs]
    return res


def map_components(
    dag_components: List[DAGCircuit], cmap_components: List[CouplingMap]
) -> Dict[int, List[int]]:
    """Returns a map where the key is the index of each connected component in cmap_components and
    the value is a list of indices from dag_components which should be placed onto it."""
    free_qubits = {index: len(cmap.graph) for index, cmap in enumerate(cmap_components)}
    out_mapping = defaultdict(list)

    for dag_index, dag in sorted(
        enumerate(dag_components), key=lambda x: x[1].num_qubits(), reverse=True
    ):
        for cmap_index in sorted(
            range(len(cmap_components)), key=lambda index: free_qubits[index], reverse=True
        ):
            # TODO: Improve heuristic to involve connectivity and estimate
            # swap cost
            if dag.num_qubits() <= free_qubits[cmap_index]:
                out_mapping[cmap_index].append(dag_index)
                free_qubits[cmap_index] -= dag.num_qubits()
                break
        else:
            raise TranspilerError(
                "A connected component of the DAGCircuit is too large for any of the connected "
                "components in the coupling map."
            )
    return out_mapping


def split_barriers(dag: DAGCircuit):
    """Mutate an input dag to split barriers into single qubit barriers."""
    for node in dag.op_nodes(Barrier):
        num_qubits = len(node.qargs)
        if num_qubits == 1:
            continue
        barrier_uuid = uuid.uuid4()
        split_dag = DAGCircuit()
        split_dag.add_qubits([Qubit() for _ in range(num_qubits)])
        for i in range(num_qubits):
            split_dag.apply_operation_back(
                Barrier(1, label=barrier_uuid), qargs=[split_dag.qubits[i]]
            )
        dag.substitute_node_with_dag(node, split_dag)


def combine_barriers(dag: DAGCircuit, retain_uuid: bool = True):
    """Mutate input dag to combine barriers with UUID labels into a single barrier."""
    qubit_indices = {bit: index for index, bit in enumerate(dag.qubits)}
    uuid_map: dict[uuid.UUID, DAGOpNode] = {}
    for node in dag.op_nodes(Barrier):
        if isinstance(node.op.label, uuid.UUID):
            barrier_uuid = node.op.label
            if barrier_uuid in uuid_map:
                other_node = uuid_map[node.op.label]
                num_qubits = len(other_node.qargs) + len(node.qargs)
                new_op = Barrier(num_qubits, label=barrier_uuid)
                new_node = dag.replace_block_with_op([node, other_node], new_op, qubit_indices)
                uuid_map[barrier_uuid] = new_node
            else:
                uuid_map[barrier_uuid] = node
    if not retain_uuid:
        for node in dag.op_nodes(Barrier):
            if isinstance(node.op.label, uuid.UUID):
                node.op.label = None


def require_layout_isolated_to_component(
    dag: DAGCircuit, components_source: Union[Target, CouplingMap]
):
    """
    Check that the layout of the dag does not require connectivity across connected components
    in the CouplingMap

    Args:
        dag: DAGCircuit to check.
        components_source: Target to check against.

    Raises:
        TranspilerError: Chosen layout is not valid for the target disjoint connectivity.
    """
    if isinstance(components_source, Target):
        coupling_map = components_source.build_coupling_map(filter_idle_qubits=True)
    else:
        coupling_map = components_source
    qubit_indices = {bit: index for index, bit in enumerate(dag.qubits)}
    component_sets = [set(x.graph.nodes()) for x in coupling_map.connected_components()]
    for inst in dag.two_qubit_ops():
        component_index = None
        for i, component_set in enumerate(component_sets):
            if qubit_indices[inst.qargs[0]] in component_set:
                component_index = i
                break
        if qubit_indices[inst.qargs[1]] not in component_sets[component_index]:
            raise TranspilerError("Chosen layout is not valid for the target disjoint connectivity")


def separate_dag(dag: DAGCircuit) -> List[DAGCircuit]:
    """Separate a dag circuit into it's connected components."""
    # Split barriers into single qubit barriers before splitting connected components
    split_barriers(dag)
    im_graph, _, qubit_map, __ = vf2_utils.build_interaction_graph(dag)
    connected_components = rx.weakly_connected_components(im_graph)
    component_qubits = []
    for component in connected_components:
        component_qubits.append({qubit_map[x] for x in component})

    qubits = set(dag.qubits)

    decomposed_dags = []
    for dag_qubits in component_qubits:
        new_dag = dag.copy_empty_like()
        new_dag.remove_qubits(*qubits - dag_qubits)
        new_dag.global_phase = 0
        for node in dag.topological_op_nodes():
            if dag_qubits.issuperset(node.qargs):
                new_dag.apply_operation_back(node.op, node.qargs, node.cargs)
        idle_clbits = []
        for bit, node in new_dag.input_map.items():
            succ_node = next(new_dag.successors(node))
            if isinstance(succ_node, DAGOutNode) and isinstance(succ_node.wire, Clbit):
                idle_clbits.append(bit)
        new_dag.remove_clbits(*idle_clbits)
        combine_barriers(new_dag)
        decomposed_dags.append(new_dag)
    return decomposed_dags
