import pprint

import numpy as np

import qiskit
from qiskit.transpiler import StagedPassManager, PassManager
from qiskit.transpiler.passes import *
from qiskit.transpiler import CouplingMap
from qiskit.transpiler.preset_passmanagers import common
from qiskit.circuit import QuantumCircuit

# Only available starting with Qiskit 0.21.0
print(qiskit.__version__)

default_staged_pm = StagedPassManager()

print(default_staged_pm.stages)

staged_pm = StagedPassManager(stages=['layout', 'translate'])

print(staged_pm.stages)

cmap = CouplingMap.from_heavy_hex(11)
layout_stage = PassManager([VF2Layout(cmap, call_limit=int(5e9))])
layout_stage += common.generate_embed_passmanager(cmap)

translation_stage = common.generate_translation_passmanager(None, ['u', 'cx'])

staged_pm = StagedPassManager(stages=['layout', 'translate'], layout=layout_stage)
staged_pm.translate = translation_stage
staged_pm.draw()

staged_pm.pre_layout = common.generate_unroll_3q(None, ['u', 'cx'])
staged_pm.post_layout = common.generate_pre_op_passmanager(coupling_map=cmap)
staged_pm.pre_translate = PassManager(Depth())

staged_pm.translate.append(Size())

staged_pm.draw()

from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
from qiskit.providers.fake_provider import FakeGuadalupeV2
from qiskit.circuit.library import XGate, HGate, RXGate, PhaseGate, TGate, TdgGate

backend = FakeGuadalupeV2()
opt_level_1_pm = generate_preset_pass_manager(1, backend)

print(type(opt_level_1_pm))

opt_level_1_pm.draw()

circuit = QuantumCircuit(4)
circuit.h(0)
circuit.cx(0, 1)
circuit.cx(0, 1)
circuit.cx(0, 1)
circuit.cx(0, 2)
circuit.cx(0, 3)
circuit.measure_all()
circuit.draw('mpl')

opt_level_1_pm.run(circuit).draw('mpl')

backend_durations = backend.target.durations()

dd_sequence = [XGate(), XGate()]

scheduling_pm = PassManager([
    ALAPScheduleAnalysis(backend_durations),
    PadDynamicalDecoupling(backend_durations, dd_sequence),
])

inverse_gate_list = [
    HGate(),
    (RXGate(np.pi / 4), RXGate(-np.pi / 4)),
    (PhaseGate(np.pi / 4), PhaseGate(-np.pi / 4)),
    (TGate(), TdgGate()),
]

logical_opt = PassManager([
    CXCancellation(),
    InverseCancellation([HGate(), (RXGate(np.pi / 4), RXGate(-np.pi / 4))])
])

opt_level_1_pm.scheduling = scheduling_pm
opt_level_1_pm.pre_layout = logical_opt
opt_level_1_pm.draw()

opt_level_1_pm.run(circuit).draw('mpl')

from qiskit.converters import dag_to_circuit
def callback(**kwargs):
    print(kwargs['pass_'])
    intermediate_circuit = dag_to_circuit(kwargs['dag'])
    intermediate_circuit.name = 'After:' +str(kwargs['pass_'])
    display(intermediate_circuit.draw('mpl'))

opt_level_1_pm.run(circuit, callback=callback)