"""Imports for the notebook."""
_req = """This notebook is written for

qiskit-aqua==0.6.0

Your code may not execute properly.
"""
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import warnings

import qiskit
if "0.6" not in qiskit.__qiskit_version__.get("qiskit-aqua"):
    warnings.warn(_req)

"""Optional: Turn off warnings for the notebook."""
warnings.filterwarnings("ignore")

"""Specific imports for QAOA with MaxCut."""
# Import the QAOA object
from qiskit.aqua.algorithms.adaptive import QAOA

# Import tools for the MaxCut problem
from qiskit.aqua.translators.ising.max_cut import (
    get_max_cut_qubitops, max_cut_value, random_graph
)
from qiskit.aqua.operators.weighted_pauli_operator import (
    Pauli, WeightedPauliOperator
)

# Import optimizers in Qiskit for finding the best parameters in the QAOA circuit
from qiskit.aqua.components.optimizers import ADAM, AQGD, COBYLA, POWELL

"""Helper function for drawing weighted graphs.

You don't need to know how this function works. You will see how to use it below.
"""
def draw_weighted(graph: nx.Graph,
                  pos_color: str = "blue",
                  neg_color: str = "red",
                  scale: float = 2.0,
                  **kwargs) -> None:
    """Shows a visual of a graph with edges scaled by weight and colored by sign.
    
    Args:
        graph: The weighted graph to visualize.
        pos_color: Color for edges with a positive weight.
        neg_color: Color for edges with a negative weight.
        scale: Floating point value to scale edge weights by
               in the visualization.
            
    Keyword Args:
        cut (List[Int]): A list of 0, 1 values specifying which
                         nodes are in which class. The number of
                         values must be equal to the number of
                         nodes in the graph.
    """
    pos = nx.spring_layout(graph)
    if "cut" in kwargs.keys():
        keys = kwargs["cut"]
        if len(keys) != len(graph.nodes):
            raise ValueError(
                f"ecolor_key has length {len(keys)} but graph has {len(graph.nodes)} nodes."
            )
        nx.draw_networkx_nodes(graph, pos, node_size=700, node_color=keys, cmap=plt.cm.Greens)
    else:
        nx.draw_networkx_nodes(graph, pos, node_size=700)
    
    col = lambda sgn: pos_color if sgn > 0 else neg_color
    
    for edge in graph.edges:
        weight = graph.get_edge_data(*edge)["weight"]
        sgn = np.sign(weight)
        size = abs(weight)
        nx.draw_networkx_edges(graph, 
                               pos, 
                               edgelist=[edge], 
                               width=scale * size,
                               edge_color=col(sgn),
                               alpha=0.5)
    nx.draw_networkx_labels(graph, pos, font_size=20)
    plt.axis("off")
    plt.show()

"""Define the graph for MaxCut via an adjacency matrix."""
nodes = 6  # Vary the number of nodes here
matrix = random_graph(n=nodes, edge_prob=0.5, seed=2)
print("The adjacency matrix is:")
print(matrix)

"""Convert the adjacency matrix to a (weighted) graph and visualize it."""
graph = nx.from_numpy_array(matrix, parallel_edges=False)
draw_weighted(graph)

"""TODO: Put a list of 1s and 0s to assign vertex sets and see
what the value of your cut is.
"""
cut = np.array()  ### <-- Your code here!
print("The value of this cut is:", max_cut_value(cut, matrix))

"""Visualize the cut by coloring nodes in distinct vertex sets different colors."""
draw_weighted(graph, cut=[0, 0, 0, 1, 1, 1])

"""Pauli operators from matrix."""
op, shift = get_max_cut_qubitops(matrix)

"""Inspect the Pauli operators."""
print("Edge set of the graph:")
for edge in graph.edges:
    print("Weight:", graph.get_edge_data(*edge)["weight"], "Edge:", edge)
print("\nWeighted Pauli operators.")
for pauli in op.paulis:
    print(2 * np.real(pauli[0]), "*", pauli[1].to_label()[::-1])

"""Make the QAOA instance."""
qaoa = QAOA(op, POWELL(), p=1)

"""See the settings of the QAOA object."""
print(qaoa.print_settings())

"""Inspect the circuits."""
backend = qiskit.BasicAer.get_backend("qasm_simulator")
circs = qaoa.construct_circuit([1, 2], backend=backend)

print(f"There are {len(circs)} circuits.")
print(circs[0])

"""Set the number of points N to define the grid. Larger N ==> longer runtime."""
N = 10
gammas = np.linspace(-np.pi, np.pi, N)
betas = np.linspace(-np.pi, np.pi, N)

"""Minor hacks for the QAOA instance to make the grid search possible.
Run this cell without too much thought -- this is necessary because the way Aqua is set up.
"""
quantum_instance = qiskit.aqua.QuantumInstance(backend=qiskit.BasicAer.get_backend("qasm_simulator"))
qaoa._quantum_instance = quantum_instance
qaoa._use_simulator_operator_mode = True

"""Do the grid search and display the progress."""
import progressbar
bar = progressbar.ProgressBar(maxval=N**2)

costs = np.zeros((len(gammas), len(betas)), dtype=float)
bar.start()
for (ii, gamma) in enumerate(gammas):
    for (jj, beta) in enumerate(betas):
        costs[ii][jj] = qaoa._energy_evaluation(np.array([gamma, beta]))
        bar.update(N * ii + jj)
bar.finish()

"""Visualize the landscape."""
plt.figure(figsize=(7, 7));
plt.imshow(costs, origin=(0, 0));
plt.xlabel("Gammas")
plt.ylabel("Betas")
plt.colorbar();

"""Your code here!"""


"""Write code to answer the above question here."""
print("Min cost =", np.min(costs))  # Your code here!

"""Get a quantum instance and run the algorithm."""
qaoa._optimizer = POWELL()
result = qaoa.run(quantum_instance)

"""View the optimal cost."""
qaoa.get_optimal_cost()

"""Get the circuit with optimal parameters."""
circ = qaoa.get_optimal_circuit()
qreg = circ.qregs[0]
creg = qiskit.ClassicalRegister(6)
circ.add_register(creg)
circ.measure(qreg, creg)
print(circ)

"""Execute the circuit to sample from it."""
job = qiskit.execute(circ, backend=backend, shots=100000)
res = job.result()
counts = res.get_counts()

"""Visualize the statistics."""
qiskit.visualization.plot_histogram(counts, figsize=(17, 6))

"""Get the top sampled bit strings."""
import operator
ntop = 10
top_cuts = sorted(counts, key=operator.itemgetter(1))
print(f"Top {ntop} sampled cuts.")
for cut in top_cuts[:ntop]:
    print(cut[::-1])

"""Select a sampled cut and see its value."""
cut = np.array()  ### <-- Your answer here!
print("The value of this cut is:", max_cut_value(cut, matrix))

"""Brute force search for the maximum cut."""
import itertools

high = -np.inf
conf = np.zeros(nodes)

cuts = itertools.product(*[[0, 1]] * nodes)

for cut in cuts:
    cur = max_cut_value(np.array(cut), matrix)
    if cur > high:
        conf = np.array(cut)
        high = cur

print("Value of maximum cut:", high)
print("Optimal cut:", conf)

"""Visualize the graph to see the maximum cut."""
draw_weighted(graph, cut=conf)