from qiskit import QuantumCircuit, QuantumRegister
import numpy as np

# load a dummy circuit
# generate network of random cnots
n, M = 2, 4
N = n**2
register = QuantumRegister(N)
qc = QuantumCircuit(register)

for _ in range(M):
    qc.cx(*np.random.choice(N, 2, replace=False))
qc.name = "random cx circuit"
qc.decompose().draw("mpl")

from qiskit.transpiler.coupling import CouplingMap

qc = QuantumCircuit(3)
qc.ccx(0, 1, 2)
coupling = CouplingMap.from_line(3)
qc.decompose().draw("mpl")

from virtual_swap.pass_managers import BruteCNS

brute = BruteCNS(coupling)
transp = brute.run(qc)
transp.draw("mpl")

from qiskit.transpiler.passes import OptimizeSwapBeforeMeasure, Unroller
from qiskit.converters import dag_to_circuit
from qiskit.transpiler import PassManager

clean = PassManager()
clean.append(OptimizeSwapBeforeMeasure())
clean.append(Unroller(["u", "cx", "iswap", "swap"]))

property_set = brute.pm.property_set  # Reference for brevity
dags, costs, perms = (
    property_set["all_dags"],
    property_set["all_costs"],
    property_set["all_perms"],
)

for dag, cost, perm in zip(dags, costs, perms):
    print(cost)
    perm_str = ", ".join(
        [f"{node.name}[{node.qargs[0].index},{node.qargs[1].index}]" for node in perm]
    )
    print("Permutation:", perm_str)
    circuit = dag_to_circuit(dag)
    print(clean.run(circuit).draw(fold=-1))