from networkx import barabasi_albert_graph, draw
graph = barabasi_albert_graph(n=10, m=6, seed=42)

draw(graph, with_labels=True)

from qopt_best_practices.utils import build_max_cut_paulis
from qiskit.quantum_info import SparsePauliOp

local_correlators = build_max_cut_paulis(graph)
cost_operator = SparsePauliOp.from_list(local_correlators)
print(cost_operator)

# import json

# graph_file = "data/graph_2layers_0seed.json"
# data = json.load(open(graph_file, "r"))

# local_correlators = data["paulis"]
# print(local_correlators)
# cost_operator = SparsePauliOp.from_list(local_correlators)

num_qubits = cost_operator.num_qubits
print(num_qubits)

qaoa_layers = 3

from qiskit.providers.fake_provider import GenericBackendV2
from qiskit.transpiler import CouplingMap

cmap = CouplingMap.from_heavy_hex(distance=3)
print(cmap.size())
backend = GenericBackendV2(num_qubits = 19, coupling_map = cmap, basis_gates = ["x", "sx", "cz", "id", "rz"], seed=0)

from qiskit import QuantumCircuit
from qiskit.circuit.library import QAOAAnsatz
from qiskit.circuit import ParameterVector

# Initial state = equal superposition
initial_state = QuantumCircuit(num_qubits)
initial_state.h(range(num_qubits))

# Mixer operator = rx rotations
betas = ParameterVector("β", qaoa_layers)
mixer_operator = QuantumCircuit(num_qubits)
mixer_operator.rx(-2*betas[0], range(num_qubits))

# Use off-the-shelf qiskit QAOAAnsatz
qaoa_ansatz = QAOAAnsatz(
    cost_operator,
    initial_state = initial_state,
    mixer_operator = mixer_operator,
    reps = qaoa_layers,
)
qaoa_ansatz.measure_all()

from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager

naive_pm = generate_preset_pass_manager(backend=backend, optimization_level=3)

print(naive_pm.expanded_stages)

# naive_pm.draw()

print("init")
for task in naive_pm.init._tasks:
    print(task)

print("----")
print("layout")
for task in naive_pm.layout._tasks:
    print(task)

print("----")
print("routing")
for task in naive_pm.routing._tasks:
    print(task)

print("----")
print("translation")
for task in naive_pm.translation._tasks:
    print(task)

print("----")
print("optimization")
for task in naive_pm.optimization._tasks:
    print(task)

print("----")
print("scheduling")
for task in naive_pm.scheduling._tasks:
    print(task)

def callback_func(**kwargs):
    pass_ = kwargs['pass_']
    print(pass_)

import time
from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager

naive_pm = generate_preset_pass_manager(backend=backend, optimization_level=3)
t0 = time.time()
naively_transpiled_qaoa = naive_pm.run(qaoa_ansatz)
t1 = time.time()

print(f"transpilation time: {t1 - t0} (s)")

naively_transpiled_qaoa.count_ops()

naively_transpiled_qaoa.depth()

# We are taking advantage of the QAOAAnsatz class to build the cost layer, 
# note that we are giving it dummy initial state and mixer circuits.

dummy_initial_state = QuantumCircuit(num_qubits)  # the real initial state is defined later
dummy_mixer_operator = QuantumCircuit(num_qubits)  # the real mixer is defined later

cost_layer = QAOAAnsatz(
    cost_operator,
    reps=1,
    initial_state=dummy_initial_state,
    mixer_operator=dummy_mixer_operator,
    name="QAOA cost block",
)

cost_layer.draw("mpl", fold=-1)

cost_layer.decompose(reps=2).draw("mpl", fold=-1)

from qopt_best_practices.qubit_selection import BackendEvaluator

# The backend evaluator finds the line of qubits with the best fidelity to map the circuit to
path_finder = BackendEvaluator(backend)
path, fidelity, num_subsets = path_finder.evaluate(num_qubits)
print(path, fidelity)

from qiskit.transpiler import PassManager
from qiskit.transpiler.passes import (
    BasisTranslator,
    UnrollCustomDefinitions,
    CommutativeCancellation,
    Decompose,
    CXCancellation,
    HighLevelSynthesis,
    InverseCancellation
)

from qiskit.transpiler.passes.routing.commuting_2q_gate_routing import (
    SwapStrategy,
    FindCommutingPauliEvolutions,
    Commuting2qGateRouter,
)
from qiskit.circuit.library.standard_gates.equivalence_library import _sel
from qiskit.circuit.library import CXGate

# 1. choose swap strategy (in this case -> line)
swap_strategy = SwapStrategy.from_line([i for i in range(num_qubits)])
edge_coloring = {(idx, idx + 1): (idx + 1) % 2 for idx in range(num_qubits)}

# 2. define pass manager for cost layer
pre_init = PassManager(
            [HighLevelSynthesis(basis_gates=['PauliEvolution']),
             FindCommutingPauliEvolutions(),
             Commuting2qGateRouter(
                    swap_strategy,
                    edge_coloring,
                ),
             HighLevelSynthesis(basis_gates=["x", "cx", "sx", "rz", "id"]),
             InverseCancellation(gates_to_cancel=[CXGate()]),
            ]
)

# Before:
print(cost_layer.decompose(reps=4).count_ops())
cost_layer.decompose(reps=4).draw("mpl", fold=-1)

tmp = pre_init.run(cost_layer)

# After:
print(tmp.count_ops())
tmp.draw('mpl', fold=-1)

from qiskit.transpiler.basepasses import TransformationPass
from qiskit.converters import circuit_to_dag, dag_to_circuit

class QAOAPass(TransformationPass):

    def __init__(self, num_layers, num_qubits, init_state = None, mixer_layer = None):

        super().__init__()
        self.num_layers = num_layers
        self.num_qubits = num_qubits
        
        if init_state is None:
            # Add default initial state -> equal superposition
            self.init_state = QuantumCircuit(num_qubits)
            self.init_state.h(range(num_qubits))
        else: 
            self.init_state = init_state
        
        if mixer_layer is None:
            # Define default mixer layer
            self.mixer_layer = QuantumCircuit(num_qubits)
            self.mixer_layer.rx(-2*betas[0], range(num_qubits))
        else:
            self.mixer_layer = mixer_layer

    def run(self, cost_layer_dag):

        cost_layer = dag_to_circuit(cost_layer_dag)
        qaoa_circuit = QuantumCircuit(self.num_qubits, self.num_qubits)
        # Re-parametrize the circuit
        gammas = ParameterVector("γ", self.num_layers)
        betas = ParameterVector("β", self.num_layers)

        # Add initial state
        qaoa_circuit.compose(self.init_state, inplace = True)

        # iterate over number of qaoa layers
        # and alternate cost/reversed cost and mixer
        for layer in range(self.num_layers): 
        
            bind_dict = {cost_layer.parameters[0]: gammas[layer]}
            bound_cost_layer = cost_layer.assign_parameters(bind_dict)
            
            bind_dict = {self.mixer_layer.parameters[0]: betas[layer]}
            bound_mixer_layer = self.mixer_layer.assign_parameters(bind_dict)
        
            if layer % 2 == 0:
                # even layer -> append cost
                qaoa_circuit.compose(bound_cost_layer, range(self.num_qubits), inplace=True)
            else:
                # odd layer -> append reversed cost
                qaoa_circuit.compose(bound_cost_layer.reverse_ops(), range(self.num_qubits), inplace=True)
        
            # the mixer layer is not reversed
            qaoa_circuit.compose(bound_mixer_layer, range(self.num_qubits), inplace=True)
        

        if self.num_layers % 2 == 1:
            # iterate over layout permutations to recover measurements
            if self.property_set["virtual_permutation_layout"]:
                for cidx, qidx in self.property_set["virtual_permutation_layout"].get_physical_bits().items():
                    qaoa_circuit.measure(qidx, cidx)
            else:
                print("layout not found, assigining trivial layout")
                for idx in range(self.num_qubits):
                    qaoa_circuit.measure(idx, idx)
        else:
            for idx in range(self.num_qubits):
                qaoa_circuit.measure(idx, idx)
    
        return circuit_to_dag(qaoa_circuit)
        

init = PassManager([QAOAPass(num_layers=3, num_qubits=10)])

tmp_out = init.run(tmp)
tmp_out.count_ops()
tmp_out.draw('mpl', fold=-1)

from qiskit.transpiler import Layout

# We use the obtained path to define the initial layout
initial_layout = Layout.from_intlist(path, cost_layer.qregs[0])

# The post init step unrolls the gates in the ansatz to the backend basis gates
post_init = PassManager(
    [
        UnrollCustomDefinitions(_sel, basis_gates=backend.operation_names, min_qubits=3),
        BasisTranslator(_sel, target_basis=backend.operation_names, min_qubits=3),
    ]
)

staged_pm = generate_preset_pass_manager(3, backend, initial_layout=initial_layout)
staged_pm.pre_init = pre_init
staged_pm.init = init
staged_pm.post_init = post_init
# staged_pm.routing = None

print("pre_init")
for task in staged_pm.pre_init._tasks:
    print(task)
    
print("----")
print("init")
for task in staged_pm.init._tasks:
    print(task)

print("----")
print("post_init")
for task in staged_pm.post_init._tasks:
    print(task)
    
print("----")
print("layout")
for task in staged_pm.layout._tasks:
    print(task)

print("----")
print("optimization")
for task in staged_pm.optimization._tasks:
    print(task)

print("----")
print("routing")
for task in staged_pm.routing._tasks:
    print(task)

t0_opt = time.time()
optimally_transpiled_qaoa = staged_pm.run(cost_layer)
t1_opt = time.time()

naive_count = naively_transpiled_qaoa.count_ops().get("cz", 0)
optimal_count = optimally_transpiled_qaoa.count_ops().get("cz", 0)

print(f"2q gate count for naive circuit = {naive_count}")
print(f"2q gate count for optimal circuit = {optimal_count}")

naive_2q_depth = naively_transpiled_qaoa.depth(filter_function=lambda x: x.operation.name == "cz")
optimal_2q_depth = optimally_transpiled_qaoa.depth(filter_function=lambda x: x.operation.name == "cz")

print(f"2q depth for naive circuit = {naive_2q_depth}")
print(f"2q depth for optimal circuit = {optimal_2q_depth}")

naive_depth = naively_transpiled_qaoa.depth()
optimal_depth = optimally_transpiled_qaoa.depth()

print(f"total depth for naive circuit = {naive_depth}")
print(f"total depth for optimal circuit = {optimal_depth}")

time_naive = t1 - t0
time_optimal = (t1_opt - t0_opt)

print(f"total time for naive transpilation = {time_naive} (s)")
print(f"total time for optimal transpilation = {time_optimal} (s)")

# optimally_transpiled_qaoa.draw(fold=-1)

# naively_transpiled_qaoa.draw(fold=-1)

from qiskit.transpiler.passes import (
    FullAncillaAllocation,
    EnlargeWithAncilla,
    ApplyLayout,
    SetLayout,
)

# The post init step unrolls the gates in the ansatz to the backend basis gates
post_init = PassManager(
    [
        UnrollCustomDefinitions(_sel, basis_gates=backend.operation_names),
        BasisTranslator(_sel, target_basis=backend.operation_names),
    ]
)

# The layout step applies the previously computed layout and enlarges the circuit 
# with ancilla qubits to have the same number of physical qubits as the target
layout = PassManager(
    [
        SetLayout(initial_layout),
        FullAncillaAllocation(backend.target),
        EnlargeWithAncilla(),
        ApplyLayout(),
    ]
)

# The optimization step performs additional gate cancellations
optimization = PassManager(
    [
      CommutativeCancellation(target=backend.target)
    ]
)

from qiskit.transpiler import StagedPassManager
other_staged_pm = StagedPassManager(stages=["init", "layout", 'optimization'], pre_init=pre_init, init=init, post_init=post_init, layout=layout, optimization=optimization)

t0_opt = time.time()
optimally_transpiled_qaoa = other_staged_pm.run(cost_layer)
t1_opt = time.time()

naive_count = naively_transpiled_qaoa.count_ops().get("cz", 0)
optimal_count = optimally_transpiled_qaoa.count_ops().get("cz", 0)

print(f"2q gate count for naive circuit = {naive_count}")
print(f"2q gate count for optimal circuit = {optimal_count}")

naive_2q_depth = naively_transpiled_qaoa.depth(filter_function=lambda x: x.operation.name == "cz")
optimal_2q_depth = optimally_transpiled_qaoa.depth(filter_function=lambda x: x.operation.name == "cz")

print(f"2q depth for naive circuit = {naive_2q_depth}")
print(f"2q depth for optimal circuit = {optimal_2q_depth}")

naive_depth = naively_transpiled_qaoa.depth()
optimal_depth = optimally_transpiled_qaoa.depth()

print(f"total depth for naive circuit = {naive_depth}")
print(f"total depth for optimal circuit = {optimal_depth}")

time_naive = t1 - t0
time_optimal = (t1_opt - t0_opt)

print(f"total time for naive transpilation = {time_naive} (s)")
print(f"total time for optimal transpilation = {time_optimal} (s)")