from virtual_swap.deprecated.cns_sa import VSwapPass

from qiskit import QuantumCircuit
from qiskit.quantum_info import Operator

qc = QuantumCircuit(3)
qc.cx(0, 1)
qc.h(1)
qc.iswap(1, 2)
qc.y(2)
original_op = Operator(qc)
qc.draw("mpl")

from qiskit.converters import circuit_to_dag

dag = circuit_to_dag(qc)
dag.draw()

# manually select the first cx node
from qiskit.circuit.library.standard_gates import CXGate, iSwapGate

h_node = dag.op_nodes(iSwapGate)[0]

import numpy as np

# cx -> iswap
cx_replace = QuantumCircuit(2, 0)
cx_replace.h(1)
cx_replace.rz(-np.pi / 2, 0)
cx_replace.rz(-np.pi / 2, 1)
cx_replace.iswap(0, 1)
cx_replace.h(0)
cx_replace.draw("mpl")

# iswap -> cx
iswap_replace = QuantumCircuit(2, 0)
iswap_replace.rz(np.pi / 2, 0)
iswap_replace.rz(np.pi / 2, 1)
iswap_replace.h(1)
iswap_replace.cx(0, 1)
iswap_replace.h(1)

from qiskit.dagcircuit import DAGOpNode, DAGInNode, DAGOutNode, DAGCircuit
from qiskit.circuit.library import iSwapGate, SwapGate


def cns_transform(dag: DAGCircuit, h_node):
    """Alternative implementation, adds nodes into blank copy of dag"""
    new_dag = dag.copy_empty_like()

    flip_flag = False

    swap_wires = {
        qarg1: qarg2 for qarg1, qarg2 in zip(h_node.qargs, h_node.qargs[::-1])
    }

    for node in dag.topological_op_nodes():
        if node == h_node:
            if node.name == "cx":
                new_dag.apply_operation_back(cx_replace.to_instruction(), node.qargs)
            elif node.name == "iswap":
                new_dag.apply_operation_back(iswap_replace.to_instruction(), node.qargs)
            else:
                raise ValueError("Unsupported operation")
            flip_flag = True

        else:
            if flip_flag:
                new_dag.apply_operation_back(
                    node.op, [swap_wires.get(qarg, qarg) for qarg in node.qargs]
                )
            else:
                new_dag.apply_operation_back(node.op, node.qargs)
    # fix with a swap
    new_dag.apply_operation_back(SwapGate(), h_node.qargs)
    return new_dag

def cns_transform_old(dag, h_node):
    swap_wires = {
        qarg1: qarg2 for qarg1, qarg2 in zip(h_node.qargs, h_node.qargs[::-1])
    }

    # NOTE, include the node itself, because we want to reference set of all outgoing edges
    node_list = list(dag.descendants(h_node))
    node_list.insert(0, h_node)

    for source in node_list:
        # first, update edge arguments
        for target in dag.successors(source):
            edge = dag._multi_graph.get_edge_data(source._node_id, target._node_id)
            dag._multi_graph.update_edge(
                source._node_id, target._node_id, swap_wires.get(edge, edge)
            )

        # second, update the node's qargs
        if source == h_node:
            # here replace the operation...
            if source.name == "cx":
                dag.substitute_node(source, cx_replace.to_instruction())
            elif source.name == "iswap":
                dag.substitute_node(source, iswap_replace.to_instruction())
            else:
                raise ValueError("Unsupported operation")

        # third, update output nodes
        elif isinstance(source, DAGOpNode):
            source.qargs = [swap_wires.get(qarg, qarg) for qarg in source.qargs]
        elif isinstance(source, DAGOutNode):
            source.wire = swap_wires.get(source.wire, source.wire)

    # add a swap a swap at end of circuit
    # NOTE, preserves unitary for equivalence checks
    # final SWAP gates can be removed by the optimizer
    dag.apply_operation_back(SwapGate(), h_node.qargs)

dag = cns_transform(dag, h_node)
# cns_transform_old(dag, h_node)
dag.draw()

from qiskit.converters import dag_to_circuit

qc = dag_to_circuit(dag)

# transpile into iswap, cx, u basis
from qiskit.transpiler import PassManager
from qiskit.transpiler.passes import Unroller

pm = PassManager([Unroller(["swap", "iswap", "cx", "u"])])
qc = pm.run(qc)

qc.draw("mpl")

# verify circuits are equivalent
from qiskit.quantum_info import Operator

original_op.equiv(Operator(qc))

# get first cx node from qc
dag = circuit_to_dag(qc)
cx_node = dag.op_nodes(CXGate)
from virtual_swap.cns_transform import cns_transform

new_dag = cns_transform(dag, *cx_node, preserve_layout=True)
new_qc = dag_to_circuit(new_dag)
new_qc.draw("mpl")