from qiskit import QuantumCircuit, transpile, schedule, IBMQ
from qiskit.transpiler import PassManager

import numpy as np

IBMQ.load_account()
provider = IBMQ.get_provider(hub='ibm-q-internal', group='deployed', project='default')

backend = provider.get_backend('ibm_lagos')

%matplotlib inline
%config InlineBackend.figure_format = 'retina'
from IPython.display import display


from qiskit.circuit.library.standard_gates.equivalence_library import StandardEquivalenceLibrary as std_eqlib

# Transpiler passes
from qiskit.transpiler.passes import Collect2qBlocks
from qiskit.transpiler.passes import ConsolidateBlocks
from qiskit.transpiler.passes import Optimize1qGatesDecomposition
from qiskit.transpiler.passes.basis import BasisTranslator, UnrollCustomDefinitions
from qiskit.transpiler.passes.calibration.builders import RZXCalibrationBuilderNoEcho

# New transpiler pass
from qiskit.transpiler.passes.optimization.echo_rzx_weyl_decomposition import EchoRZXWeylDecomposition

from qiskit.circuit.library.standard_gates import (RZZGate, RXXGate, RYYGate, RZXGate, 
                                                   CPhaseGate, CRZGate, HGate, SwapGate, CXGate, iSwapGate)

import qiskit.quantum_info as qi
from qiskit.quantum_info.synthesis.two_qubit_decompose import (TwoQubitControlledUDecomposer, 
                                                               TwoQubitWeylDecomposition)

theta = 0.8

qc = QuantumCircuit(2)

qc.cx(0, 1)
qc.rz(theta, 1)
qc.cx(0, 1)

qc.draw('mpl')

qct = transpile(qc, backend)
schedule(qct, backend).draw()

theta = 0.8

# equivalent decomposition in terms of echoed Cross-Resonance gates.

qc = QuantumCircuit(2)

qc.rz(np.pi / 2, 1)
qc.sx(1)
qc.rz(np.pi / 2, 1)
qc.rzx(theta / 2, 0, 1)
qc.x(0)
qc.rzx(-theta / 2, 0, 1)
qc.x(0)
qc.rz(np.pi / 2, 1)
qc.sx(1)
qc.rz(np.pi / 2, 1)

qc.draw('mpl')

# Add the calibrations.
pm = PassManager([RZXCalibrationBuilderNoEcho(backend)])

qc_pulse_efficient = pm.run(qc)

# qc_pulse_efficient.calibrations

schedule(qc_pulse_efficient, backend).draw()

# Compare the schedule durations
print('Duration of standard CNOT-based circuit:')
print(schedule(qct, backend).duration)
print('Duration of pulse-efficient circuit:')
print(schedule(qc_pulse_efficient, backend).duration)





# random two-qubit circuit

qc = QuantumCircuit(2)
qc.rzx(0.2, 1, 0)
qc.cp(1, 0, 1)
qc.s(0)
qc.rzz(-0.6, 0, 1)
qc.swap(0, 1)
qc.h(0)
qc.ryy(2.3, 0, 1)
qc.iswap(0, 1)

qc.draw('mpl')

unitary = qi.Operator(qc)

decomposer_weyl = TwoQubitWeylDecomposition(unitary)
decomposer_weyl.circuit().draw('mpl')

gate = RZXGate

two_qubit_decompose = TwoQubitControlledUDecomposer(gate)
new_circuit = two_qubit_decompose(unitary)

new_circuit.draw('mpl')

# Random circuit

delta = 8 * np.pi / 5
epsilon = np.pi / 2
eta = -5.1
theta = 0.02

qc = QuantumCircuit(3)
qc.ryy(theta, 0, 1)
qc.s(0)
qc.rz(eta, 1)
qc.rzz(epsilon, 0, 1)
qc.swap(0, 1)
qc.cx(1, 0)
qc.rzz(delta, 1, 2)
qc.swap(1, 2)
qc.h(2)

print('Original circuit:')
display(qc.draw('mpl', idle_wires=False))

qct = transpile(qc, backend)

print('Standard Qiskit transpilation:')
display(qct.draw('mpl', idle_wires=False))

pm = PassManager(
        [
            # Consolidate consecutive two-qubit operations.
            Collect2qBlocks(),
            ConsolidateBlocks(basis_gates=['rz', 'sx', 'x', 'rxx']),
        ]
)
    
qct_consolidated = pm.run(qct)
qct_consolidated.draw('mpl', idle_wires=None)

pm = PassManager(
        [
            # Rewrite circuit in terms of Weyl-decomposed echoed RZX gates.
            EchoRZXWeylDecomposition(backend),
        ]
)
    
qct_rzx_weyl = pm.run(qct_consolidated )
qct_rzx_weyl.draw('mpl', idle_wires=None)

# Random circuit

delta = 8 * np.pi / 5
epsilon = np.pi / 2
eta = -5.1
theta = 0.02

qc = QuantumCircuit(3)
qc.ryy(theta, 0, 1)
qc.s(0)
qc.rz(eta, 1)
qc.rzz(epsilon, 0, 1)
qc.swap(0, 1)
qc.cx(1, 0)
qc.rzz(delta, 1, 2)
qc.swap(1, 2)
qc.h(2)

print('Original circuit:')
display(qc.draw('mpl', idle_wires=False))

qct = transpile(qc, backend)

print('Standard Qiskit transpilation:')
display(qct.draw('mpl', idle_wires=False))

rzx_basis = ['rzx', 'rz', 'x', 'sx']

# Build a pass manager that contains all needed transpiler passes
pm = PassManager(
        [
            # Consolidate consecutive two-qubit operations.
            Collect2qBlocks(),
            ConsolidateBlocks(basis_gates=['rz', 'sx', 'x', 'rxx']),
            
            # Rewrite circuit in terms of Weyl-decomposed echoed RZX gates.
            EchoRZXWeylDecomposition(backend),
            
            # Attach scaled CR pulse schedules to the RZX gates.
            RZXCalibrationBuilderNoEcho(backend),
            
            # Simplify single-qubit gates.
            UnrollCustomDefinitions(std_eqlib, rzx_basis),
            BasisTranslator(std_eqlib, rzx_basis),
            Optimize1qGatesDecomposition(rzx_basis)
        ]
    )
    
# Run the pass manager    
qc_pulse_efficient = pm.run(qct)

print('Pulse-efficient, RZX-based circuit:')
display(qc_pulse_efficient.draw('mpl', idle_wires=False))

# Compare the schedule durations
print('Duration of standard CNOT-based circuit:')
print(schedule(qct, backend).duration)
print('Duration of pulse-efficient circuit:')
print(schedule(qc_pulse_efficient, backend).duration)





