%matplotlib inline

import numpy as np
import networkx as nx

from tqdm.notebook import tqdm

from docplex.mp.model import Model
from qiskit_optimization.translators import from_docplex_mp
from qrao.encoding import QuantumRandomAccessEncoding, EncodingCommutationVerifier

num_nodes = 6
elist = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 5), (5, 0), (0, 3), (1, 4), (2, 4)]
edges = np.zeros((num_nodes, num_nodes))
for i, j in elist:
    edges[i, j] = (i + 1) * (j + 2)

mod = Model("maxcut")
nodes = list(range(num_nodes))
var = [mod.binary_var(name="x" + str(i)) for i in nodes]
mod.maximize(
    mod.sum(
        edges[i, j] * (1 - (2 * var[i] - 1) * (2 * var[j] - 1))
        for i in nodes
        for j in nodes
    )
)
graph = nx.from_edgelist(elist)
nx.draw(graph)
problem = from_docplex_mp(mod)
print(problem)

encoding = QuantumRandomAccessEncoding(max_vars_per_qubit=3)
encoding.encode(problem)

print("Encoded Problem:\n=================")
print(encoding.qubit_op)  # The Hamiltonian without the offset
print("Offset = ", encoding.offset)
print("Variables encoded on each qubit: ", encoding.q2vars)

verifier = EncodingCommutationVerifier(encoding)
progress_bar = tqdm(verifier)
for str_dvars, obj_val, encoded_obj_val in progress_bar:
    if not np.isclose(encoded_obj_val, obj_val, atol=1e-4):
        progress_bar.disp(bar_style="danger")
        raise Exception(  # pylint: disable=broad-exception-raised
            f"Violation identified: {str_dvars} evaluates to {obj_val} "
            f"but the encoded problem evaluates to {encoded_obj_val}."
        )

import qiskit.tools.jupyter  # pylint: disable=unused-import,wrong-import-order

%qiskit_version_table
%qiskit_copyright