from src.clonk.backend_utils.topology_visualization import pretty_print
from src.clonk.backend_utils.mock_backends import FakeModular

pb = FakeModular(module_size=5, children=4, total_levels=2)
pretty_print(pb)

from qiskit.quantum_info.random import random_unitary
from qiskit import QuantumCircuit
from src.clonk.utils.riswap_gates.riswap import RiSwapGate
from src.clonk.utils.transpiler_passes.weyl_decompose import RootiSwapWeylDecomposition
from qiskit.transpiler.passes import CountOps
from qiskit.transpiler import PassManager
from tqdm import tqdm

N = 2000
basis_gate = RiSwapGate(0.5)

pm0 = PassManager()
pm0.append(RootiSwapWeylDecomposition(basis_gate=basis_gate))
pm0.append(CountOps())

res = 0
for _ in tqdm(range(N)):
    qc = QuantumCircuit(2)
    qc.append(random_unitary(dims=4), [0, 1])
    pm0.run(qc)
    res += pm0.property_set["count_ops"]["riswap"]
print("Haar score:", res / N)

from src.clonk.benchmark_suite.backend_suite_v3 import simple_backends_v3

print([backend.label for backend in simple_backends_v3])

from src.clonk.benchmark_suite.circuit_suite import circuits

q_size = 4
circuits["QAOA_Vanilla"].circuit_lambda(q_size).decompose().draw()

"""Example:"""

from src.clonk.benchmark_suite.main_plotting import benchmark, plot_wrap

for circuit_gen in circuits.values():
    benchmark(
        backends=simple_backends_v3,
        circuit_generator=circuit_gen,
        q_range=[4, 6, 8, 12, 14, 16],
        continuously_save=1,
        overwrite=0,  # NOTE: turn this to 1 if you want to scrap the saved data and recollect a new batch
        repeat=1,
    )

# NOTE when plotting use motivation = 1 to plot SWAP counts, and motivation = 0 to plot gate durations
plot_wrap(simple_backends_v3, circuits.keys(), motivation=True, plot_average=True)

"""Fig 4"""
from src.clonk.benchmark_suite.backend_suite_v2 import motivation_backends

for circuit_gen in circuits.values():
    benchmark(
        backends=motivation_backends,
        circuit_generator=circuit_gen,
        q_range=motivation_backends[0].q_range,
        continuously_save=True,
        overwrite=False,
        repeat=1,
    )
plot_wrap(motivation_backends, circuits.keys(), motivation=True, plot_average=True)

"""Fig 10"""
from src.clonk.benchmark_suite.backend_suite_v2 import small_results_backends

for circuit_gen in circuits.values():
    benchmark(
        backends=small_results_backends,
        circuit_generator=circuit_gen,
        q_range=small_results_backends[0].q_range,
        continuously_save=True,
        overwrite=False,
        repeat=1,
    )
plot_wrap(small_results_backends, circuits.keys(), motivation=True, plot_average=True)

# """Fig 12"""
from src.clonk.benchmark_suite.backend_suite_v2 import results_backends

for circuit_gen in circuits.values():
    benchmark(
        backends=results_backends,
        circuit_generator=circuit_gen,
        q_range=results_backends[0].q_range,
        continuously_save=True,
        overwrite=False,
        repeat=1,
    )
plot_wrap(results_backends, circuits.keys(), motivation=True, plot_average=True)

"""Fig 13"""
from src.clonk.benchmark_suite.backend_suite_v2 import small_results_part2_backends

for circuit_gen in circuits.values():
    benchmark(
        backends=small_results_part2_backends,
        circuit_generator=circuit_gen,
        q_range=small_results_part2_backends[0].q_range,
        continuously_save=True,
        overwrite=False,
        repeat=1,
    )
plot_wrap(
    small_results_part2_backends, circuits.keys(), motivation=False, plot_average=True
)

"""Fig 14"""
plot_wrap(results_backends, circuits.keys(), motivation=False, plot_average=True)

from src.clonk.benchmark_suite.backend_suite_v2 import small_results_part2_backendsv2
from qiskit.converters import circuit_to_dag
import numpy as np

ignore = ["u"]

for circuit_gen in circuits.values():  # [circuits['Quantum_Volume']]:
    print(circuit_gen.label)
    qc = circuit_gen.circuit_lambda(16)

    for backend in small_results_part2_backendsv2:
        print(backend.label)
        c = backend.pass_manager.run(qc)  # transpile :)
        d = circuit_to_dag(c)
        w = d.qubits  # if use wires error bc returns classical bits

        qubit_wire_counts = np.zeros(20)
        for i, wi in enumerate(w):
            for node in d.nodes_on_wire(wi, only_ops=True):
                if node.name in ignore:
                    continue
                # count the 2Q ops
                if node.name in ["cx", "fSim", "riswap"]:
                    qubit_wire_counts[i] += 1

        # print(qubit_wire_counts)
        print(sum(qubit_wire_counts))

from src.clonk.benchmark_suite.nuop_script import create_plot2, collect_random2q_data

N = 20
base_fidelity_list = [0.97, 0.98, 1 - 10e-3, 1 - 5e-3, 1 - 10e-4, 1]
filename = f"src/clonk/benchmark_suite/data-archive2/data1_random.h5"  # NOTE preloaded, change name of file to recollect
gate_error, decomp_error, fidelity_error = collect_random2q_data(
    1 - 10e-3, N=N, mode="random", fn=filename
)

create_plot2(gate_error, decomp_error, fidelity_error, plot_bool=0, fn=filename);

create_plot2(gate_error, decomp_error, fidelity_error, plot_bool=1, fn=filename);