import json
import networkx as nx

graph_file = "data/graph_2layers_0seed.json"
data = json.load(open(graph_file, "r"))
graph = nx.from_edgelist(data["Original graph"])
nx.draw(graph, with_labels=True)

num_qubits = len(graph.nodes)
print(num_qubits)

from qopt_best_practices.utils import build_max_cut_paulis, build_max_cut_graph

original_hamiltonian = build_max_cut_paulis(graph)
print(original_hamiltonian)

original_graph = build_max_cut_graph(original_hamiltonian)
nx.draw(original_graph, with_labels=True)

from qiskit.transpiler.passes.routing.commuting_2q_gate_routing import SwapStrategy

swap_strategy = SwapStrategy.from_line([i for i in range(num_qubits)])

from qopt_best_practices.sat_mapping import SATMapper

sm = SATMapper()
remapped_g, sat_map, min_sat_layers = sm.remap_graph_with_sat(
    graph=graph, swap_strategy=swap_strategy
)

print("Map from old to new edges: ", sat_map)
print("Min SAT layers:", min_sat_layers)
nx.draw(remapped_g, with_labels=True)

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

pauli_list = build_max_cut_paulis(remapped_g)
print(pauli_list)

# define a qiskit SparsePauliOp from the list of paulis
qaoa_hamiltonian = SparsePauliOp.from_list(pauli_list)
print(qaoa_hamiltonian)

from qopt_best_practices.swap_strategies import create_qaoa_swap_circuit

theta = [1, 1, 0, 1]
# we define the edge_coloring map so that RZZGates are positioned next to SWAP gates to exploit CX cancellations
edge_coloring = {(idx, idx + 1): (idx + 1) % 2 for idx in range(qaoa_hamiltonian.num_qubits)}

qaoa_circ = create_qaoa_swap_circuit(qaoa_hamiltonian, swap_strategy, edge_coloring, qaoa_layers=2)

qaoa_circ.draw("mpl")

qaoa_circ.decompose(reps=2).draw("mpl")

# SIMULATED FAKE BACKEND --> BackendSampler

from qiskit_ibm_runtime.fake_provider import FakeGuadalupe
from qiskit.primitives import BackendSampler

backend = FakeGuadalupe()
sampler = BackendSampler(backend=backend, skip_transpilation=True)

## UNCOMMENT TO RUN ON:
## REAL BACKEND --> Sampler from qiskit_ibm_runtime

# from qiskit_ibm_runtime import Sampler, Options, QiskitRuntimeService

# service = QiskitRuntimeService(channel='ibm_quantum')
# backend = service.get_backend('Guadalupe')
# options = Options()
# options.transpiler.skip_transpilation = True

# sampler = Sampler(backend=backend, options=options)

from qopt_best_practices.qubit_selection import BackendEvaluator

path_finder = BackendEvaluator(backend)

# the Backend Evaluator accepts custom subset definitions and metrics,
# but defaults to finding the line with the best fidelity
path, fidelity, num_subsets = path_finder.evaluate(num_qubits)

print("Best path: ", path)
print("Best path fidelity", fidelity)
print("Num. evaluated paths", num_subsets)

from qiskit.transpiler import Layout

initial_layout = Layout.from_intlist(path, qaoa_circ.qregs[0])  # needs qaoa_circ

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

from qiskit import transpile

basis_gates = ["rz", "sx", "x", "cx"]

backend_cmap = CouplingMap(backend.configuration().coupling_map)

pass_manager_post = PassManager(
    [
        SetLayout(initial_layout),
        FullAncillaAllocation(backend_cmap),
        EnlargeWithAncilla(),
        ApplyLayout(),
    ]
)

# Map to initial_layout and finally enlarge with ancilla.
qaoa_circ = pass_manager_post.run(qaoa_circ)

# Now transpile to sx, rz, x, cx basis
qaoa_circ = transpile(qaoa_circ, basis_gates=basis_gates)
qaoa_circ.draw("mpl", idle_wires=False)

from qopt_best_practices.cost_function import evaluate_sparse_pauli


def cost_func_sampler(params, ansatz, hamiltonian, sampler):

    job = sampler.run(ansatz, params)
    sampler_result = job.result()
    sampled = sampler_result.quasi_dists[0]

    # a dictionary containing: {state: (measurement probability, value)}
    evaluated = {
        state: (probability, evaluate_sparse_pauli(state, hamiltonian))
        for state, probability in sampled.items()
    }

    result = sum(probability * value for probability, value in evaluated.values())

    return result

import numpy as np

# TQA initialization parameters
dt = 0.75
p = 2  # 2 qaoa layers
grid = np.arange(1, p + 1) - 0.5
init_params = np.concatenate((1 - grid * dt / p, grid * dt / p))
print(init_params)

from scipy.optimize import minimize

result = minimize(
    cost_func_sampler,
    init_params,
    args=(qaoa_circ, qaoa_hamiltonian, sampler),
    method="COBYLA",
)
print(result)

# auxiliary functions to sample most likely bitstring
def to_bitstring(integer, num_bits):
    result = np.binary_repr(integer, width=num_bits)
    return [int(digit) for digit in result]


def sample_most_likely_bitstring(state_vector, num_bits):
    values = list(state_vector.values())
    most_likely = np.argmax(np.abs(values))
    most_likely_bitstring = to_bitstring(most_likely, num_bits)
    most_likely_bitstring.reverse()
    return np.asarray(most_likely_bitstring)

qc = qaoa_circ.assign_parameters(result.x)
samp_dist = sampler.run(qc, shots=int(1e4)).result().quasi_dists[0]

best_result = sample_most_likely_bitstring(samp_dist, len(graph))
print(best_result)

import matplotlib.pyplot as plt
import networkx as nx

# auxiliary function to plot graphs
def plot_result(G, x):
    colors = ["r" if i == 0 else "b" for i in x]
    pos, default_axes = nx.spring_layout(G), plt.axes(frameon=True)
    nx.draw_networkx(G, node_color=colors, node_size=600, alpha=0.8, pos=pos)

plot_result(graph, best_result)