from qiskit import QuantumCircuit, transpile
from qiskit.opflow import PauliSumOp
from qiskit.circuit.library import PauliEvolutionGate
from qiskit.transpiler import Layout, CouplingMap, PassManager
from qiskit.transpiler.passes import FullAncillaAllocation
from qiskit.transpiler.passes import EnlargeWithAncilla
from qiskit.transpiler.passes import ApplyLayout
from qiskit.transpiler.passes import SetLayout
from qiskit.transpiler.passes.routing.commuting_2q_gate_routing import (
    SwapStrategy,
    FindCommutingPauliEvolutions,
    Commuting2qGateRouter,
)

# Define the circuit on logical qubits
op = PauliSumOp.from_list([("IZZI", 1), ("ZIIZ", 2), ("ZIZI", 3)])
circ = QuantumCircuit(4)
circ.append(PauliEvolutionGate(op, 1), range(4))
circ.draw("mpl")

PassManager([FindCommutingPauliEvolutions()]).run(circ).draw("mpl")

swap_layers = (
    ((1, 3), ),         # Layer 1: swap qubits 1 & 3
    ((0, 1), (3, 4)),   # Layer 2: simultaneouosly swap qubits 0 & 1 and 3 & 4
    ((1, 3), )          # Layer 3: swap qubits 1 & 3
)

cmap = CouplingMap(couplinglist=[(0, 1), (1, 2), (1, 3), (3, 4)])
swap_strat = SwapStrategy(cmap, swap_layers)
print(f"Missing connectivity: {swap_strat.missing_couplings}")

backend_cmap = CouplingMap(couplinglist=[(0, 1), (1, 2), (1, 3), (3, 4)])

def make_passmanager(cmap, swap_strat, initial_layout):
    """Creates the passes needed."""
    return PassManager(
    [
        FindCommutingPauliEvolutions(),
        Commuting2qGateRouter(swap_strat),
        SetLayout(initial_layout),
        FullAncillaAllocation(backend_cmap),
        EnlargeWithAncilla(),
        ApplyLayout(),
    ]
)

# Define the swap strategy on qubits before the initial_layout is applied.
swap_cmap = CouplingMap(couplinglist=[(0, 1), (1, 2), (2, 3)])
swap_strat = SwapStrategy(swap_cmap, swap_layers=[[(1, 2)], [(0, 1), (2, 3)]])
# Chose qubits 0, 1, 3, and 4 from the backend coupling map shown above.
backend_cmap = CouplingMap(couplinglist=[(0, 1), (1, 2), (1, 3), (3, 4)])
initial_layout = Layout.from_intlist([0, 1, 3, 4], *circ.qregs)

pm = make_passmanager(backend_cmap, swap_strat, initial_layout)

# Insert swap gates, map to initial_layout and finally enlarge with ancilla.
pm.run(circ).draw("mpl")

pm.run(circ).decompose().decompose().draw("mpl")

op = PauliSumOp.from_list(
    [
        ("IIIZZ", 1), ("IIZIZ", 1), ("IZIIZ", 1), ("ZIIIZ", 2), 
        ("IIZZI", 3), ("IZIZI", 3), ("ZIIZI", 4),
        ("IZZII", 1), ("ZIZII", 2),
        ("ZZIII", 2),
    ]
)
circ = QuantumCircuit(5)
circ.append(PauliEvolutionGate(op, 1), range(5))
circ.draw("mpl")

swap_cmap = CouplingMap(couplinglist=[(0, 1), (1, 2), (1, 3), (3, 4)])

swap_strat = SwapStrategy(swap_cmap, swap_layers=[[(1, 3)], [(0, 1), (3, 4)], [(1, 3)]])
swap_strat.missing_couplings

initial_layout = Layout.from_intlist([0, 1, 2, 3, 4], *circ.qregs)
pm = make_passmanager(backend_cmap, swap_strat, initial_layout)

swapped_circ = pm.run(circ)
swapped_circ.draw("mpl")

tswapped = transpile(swapped_circ, basis_gates=["rz", "cx"], optimization_level=2)
n_cx = tswapped.count_ops()["cx"]
n_xsx = tswapped.count_ops().get("sx", 0) + tswapped.count_ops().get("x", 0)
print(f"Number of CNOTs: {n_cx}\nNumber of X and SX: {n_xsx}.")
print(f"Depth {tswapped.depth()}.")
tswapped.draw("mpl", fold=False)

opt3_circ = transpile(circ, basis_gates=["rz", "cx", "sx", "x"], optimization_level=3, coupling_map=backend_cmap)
n_cx = opt3_circ.count_ops()["cx"]
n_xsx = opt3_circ.count_ops().get("sx", 0) + tswapped.count_ops().get("x", 0)
print(f"Number of CNOTs: {n_cx}\nNumber of X and SX: {n_xsx}.")
print(f"Depth {opt3_circ.depth()}.")
opt3_circ.draw("mpl")