from qiskit_tket_passes import ToQiskitPass
import pytket.passes as tkps

_pass1 = ToQiskitPass(tkps.SynthesiseTket)

# If TKET pass' constructor expects some arguments, pass them after the class name
_pass2 = ToQiskitPass(tkps.FullPeepholeOptimise, allow_swaps=True, target_2qb_gate='cx')

from pytket.architecture import Architecture

arc = Architecture([(0, 1), (1, 0), (1, 2), (1, 3), (2, 1), (3, 1), (3, 4), (4, 3)])
_pass = ToQiskitPass(tkps.AASRouting, arc=arc)

_pass = ToQiskitPass(tkps.AASRouting, arc=[[0, 1], [1, 0], [1, 2], [1, 3], [2, 1], [3, 1], [3, 4], [4, 3]])

from qiskit.providers.fake_provider import FakeQuitoV2

backend = FakeQuitoV2()
_pass = ToQiskitPass(tkps.AASRouting, target=backend.target)

print(_pass.tket_argument('arc').to_dict())

from qiskit.circuit.random import random_circuit
from pytket.extensions.qiskit import AerBackend
from qiskit_tket_passes import TketPassManager

qc = random_circuit(3, 10, seed=1)
pm = TketPassManager(AerBackend(), optimization_level=2)
transpiled_qc = pm.run(qc)

from qiskit.circuit.random import random_circuit

circ = random_circuit(7, 30, seed=1)
circ.draw('mpl', fold=-1)

from qiskit.transpiler.passes import (
    UnitarySynthesis,
    Collect2qBlocks,
    ConsolidateBlocks,
    Unroll3qOrMore,
)
from qiskit.transpiler import PassManager, StagedPassManager

basis_gates = ["rx", "ry", "rxx"]
init = PassManager([UnitarySynthesis(basis_gates, min_qubits=3), Unroll3qOrMore()])
translate = PassManager(
    [
        Collect2qBlocks(),
        ConsolidateBlocks(basis_gates=basis_gates),
        UnitarySynthesis(basis_gates),
    ]
)

staged_pm = StagedPassManager(
    stages=["init", "translation"], init=init, translation=translate
)

tr_circ = staged_pm.run(circ)
print('Depth:', tr_circ.depth())

optimize = PassManager(
    [
        ToQiskitPass(tkps.RemoveRedundancies)
    ]
)

staged_pm = StagedPassManager(
    stages=["init", "translation", "optimization"], init=init, translation=translate, optimization=optimize
)

tr_circ = staged_pm.run(circ)
print('Depth:', tr_circ.depth())

from pytket.extensions.qiskit import IBMQBackend
from qiskit_tket_passes import TketPassManager

qbackend = IBMQBackend('ibmq_jakarta')

pm = TketPassManager(qbackend, optimization_level=2)
#pm.draw()

tr_circ = pm.run(circ)
print('CNOTs:', tr_circ.count_ops()['cx'])
print('Depth:', tr_circ.depth())

from qiskit_ibm_provider import IBMProvider

provider = IBMProvider()
backend = provider.get_backend('ibmq_jakarta')

from qiskit import transpile

tr_circ = transpile(circ, backend=backend, optimization_level=3)
print('CNOTs:', tr_circ.count_ops()['cx'])
print('Depth:', tr_circ.depth())


print('init:')
circ_init = transpile(circ, backend=backend, optimization_level=3, init_method='tket')
print('CNOTs:', circ_init.count_ops()['cx'])
print('Depth:', circ_init.depth())

print('\ntranslation:')
circ_trans = transpile(circ, backend=backend, optimization_level=3, translation_method='tket')
print('CNOTs:', circ_trans.count_ops()['cx'])
print('Depth:', circ_trans.depth())

print('\noptimization:')
circ_opt = transpile(circ, backend=backend, optimization_level=3, optimization_method='tket')
print('CNOTs:', circ_opt.count_ops()['cx'])
print('Depth:', circ_opt.depth())


from colorama import Fore, Style

def _callback(pass_, dag, time, property_set, count):
    is_tket_pass = pass_.name().startswith('TketPass_')
    print(Fore.GREEN if is_tket_pass else Fore.WHITE, pass_.name())
    Style.RESET_ALL

circ_layout = transpile(circ, backend=backend, optimization_level=3,
    init_method='tket',
    translation_method='tket',
    optimization_method='tket',
    callback=_callback)