# Install the plugin
# !pip install -e .

import numpy as np
from qiskit import QuantumCircuit
from qiskit.quantum_info import random_clifford

from qiskit.transpiler import PassManager
from qiskit_transpiler_service.ai.synthesis import AICliffordSynthesis
from qiskit_transpiler_service.ai.collection import CollectCliffords

from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
from qiskit_ibm_runtime import QiskitRuntimeService

coupling_map = QiskitRuntimeService().backend("ibm_cairo").coupling_map

circuit = QuantumCircuit(27)

for c in range(3):
    nq = 8
    qs = np.random.choice(range(circuit.num_qubits), size=nq, replace=False)
    circuit.compose(random_clifford(nq).to_circuit(), qubits=qs.tolist(), inplace=True)
    for q in qs:
        circuit.t(q)

print(
    f"Original circuit -> Depth: {circuit.decompose(reps=3).depth()}, Gates(2q): {circuit.decompose(reps=3).num_nonlocal_gates()}"
)
circuit.draw(output="mpl", fold=-1, scale=0.3, style="iqp")

qiskit_lvl3_transpiler = generate_preset_pass_manager(
    optimization_level=3, coupling_map=coupling_map
)

lvl3_transpiled_circuit = qiskit_lvl3_transpiler.run(circuit)

print(
    f"Qiskit lvl3 Transpiler -> Depth: {lvl3_transpiled_circuit.decompose(reps=3).depth()}, Gates(2q): {lvl3_transpiled_circuit.decompose(reps=3).num_nonlocal_gates()}"
)
lvl3_transpiled_circuit.draw(output="mpl", fold=-1, scale=0.2, style="iqp")

ai_optimize_cliffords = PassManager(
    [
        CollectCliffords(),
        AICliffordSynthesis(backend_name="ibm_canberra"),
    ]
)

# AI Synthesis passes respect the coupling map and should run after transpiling
ai_optimized_circuit = ai_optimize_cliffords.run(lvl3_transpiled_circuit)

print(
    f"AI-Optimized circuit -> Depth: {ai_optimized_circuit.decompose(reps=3).depth()}, Gates(2q): {ai_optimized_circuit.decompose(reps=3).num_nonlocal_gates()}"
)
ai_optimized_circuit.draw(output="mpl", fold=-1, scale=0.25, style="iqp")

