# This code is part of Qiskit.
#
# (C) Copyright IBM 2021.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.


#Import essential packages
import matplotlib.pyplot as plt
import matplotlib.axes as axes
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter

%matplotlib inline
import numpy as np
import networkx as nx

from qiskit import Aer, IBMQ
from qiskit import QuantumRegister, ClassicalRegister, QuantumCircuit, transpile, assemble
from qiskit.providers.ibmq import least_busy
from qiskit.tools.monitor import job_monitor
from qiskit.visualization import plot_histogram
from qiskit.optimization.applications.ising import max_cut
from qiskit.optimization.applications.ising.common import sample_most_likely

# setup aqua logging
#import logging
#from qiskit.aqua import set_qiskit_aqua_logging
# set_qiskit_aqua_logging(logging.DEBUG)  # choose INFO, DEBUG to see the log

# Generating a random graph with n nodes 
n = 12
V = range(n)#np.arange(0,n,1)
E=[]
edge_bool=np.random.randint(2,size=n*(n-1)) #generate a sequence of boolen numbers for all possible edges 

W=np.zeros([n,n])
idx=0
for i in range(n):
    for j in range(i):
        temp=np.random.randint(2)
        if temp==1:
            E.append((i,j,1.0))
            W[i,j]=1
            W[j,i]=1
        #if edge_bool[idx]: # if the edge_bool is 1, then add an edge between i and j vertices with a weight 1.0.
        idx+=1

G = nx.Graph()# Generate a graph
G.add_nodes_from(V) # Add the nodes
G.add_weighted_edges_from(E) # Add the edges from E

# Generate plot of the Graph
colors       = ['r' for node in G.nodes()]
pos          = nx.spring_layout(G)

def draw_graph(G, colors, pos):
    default_axes = plt.axes(frameon=True)
    nx.draw_networkx(G, node_color=colors, node_size=600, alpha=1, ax=default_axes, pos=pos)
    #edge_labels = nx.get_edge_attributes(G, 'weight')
    #nx.draw_networkx_edge_labels(G, pos=pos, edge_labels=edge_labels)

draw_graph(G, colors, pos)

#nx.draw_networkx(G, node_color=colors, node_size=600, alpha=1, ax=default_axes, pos=pos)

print(W)

best_cost_brute = 0
for b in range(2**n):
    x = [int(t) for t in reversed(list(bin(b)[2:].zfill(n)))]
    cost = 0
    for i in range(n):
        for j in range(n):
            cost = cost + W[i,j]*x[i]*(1-x[j])
    if best_cost_brute < cost:
        best_cost_brute = cost
        xbest_brute = x
    #print('case = ' + str(x)+ ' cost = ' + str(cost))

colors = ['c' if xbest_brute[i] == 0 else 'r' for i in range(n)]
draw_graph(G, colors, pos)
print('\nBest solution = ' + str(xbest_brute) + ' cost = ' + str(best_cost_brute))

# Evaluate the function
step_size   = 0.001;

a_gamma         = np.arange(0, np.pi, step_size)
a_beta          = np.arange(0, np.pi, step_size)
a_gamma, a_beta = np.meshgrid(a_gamma,a_beta)

F1 = 3-(np.sin(2*a_beta)**2*np.sin(2*a_gamma)**2-0.5*np.sin(4*a_beta)*np.sin(4*a_gamma))*(1+np.cos(4*a_gamma)**2)

# Grid search for the minimizing variables
result = np.where(F1 == np.amax(F1))
a      = list(zip(result[0],result[1]))[0]

gamma  = a[0]*step_size;
beta   = a[1]*step_size;

#The smallest parameters and the expectation can be extracted
print('\n --- OPTIMAL PARAMETERS --- \n')
print('The maximal expectation value is:  M1 = %.03f' % np.amax(F1))
print('This is attained for gamma = %.03f and beta = %.03f' % (gamma,beta))


# prepare the quantum and classical resisters
QAOA = QuantumCircuit(len(V), len(V))

# apply the layer of Hadamard gates to all qubits
QAOA.h(range(len(V)))
QAOA.barrier()

# apply the Ising type gates with angle gamma along the edges in E
for edge in E:
    k = edge[0]
    l = edge[1]
    QAOA.cp(-2*gamma, k, l)
    QAOA.p(gamma, k)
    QAOA.p(gamma, l)
    
# then apply the single qubit X rotations with angle beta to all qubits
QAOA.barrier()
QAOA.rx(2*beta, range(len(V)))

# Finally measure the result in the computational basis
QAOA.barrier()
QAOA.measure(range(len(V)),range(len(V)))

### draw the circuit for comparison
QAOA.draw(output='mpl')

# Compute the value of the cost function
def cost_function_C(x,G):
    
    E = G.edges()
    if( len(x) != len(G.nodes())):
        return np.nan
        
    C = 0;
    for index in E:
        e1 = index[0]
        e2 = index[1]
        
        w=1.0
        C = C + w*x[e1]*(1-x[e2]) + w*x[e2]*(1-x[e1])
        
    return C

# run on local simulator
backend = Aer.get_backend("qasm_simulator")
shots = 200000

TQAOA = transpile(QAOA, backend)
qobj = assemble(TQAOA)
QAOA_results = backend.run(qobj).result()

plot_histogram(QAOA_results.get_counts(),figsize = (12,6),bar_labels = False)


# Evaluate the data from the simulator
counts = QAOA_results.get_counts()

avr_C       = 0
max_C       = [0,0]
hist        = {}

for k in range(len(G.edges())+1):
    hist[str(k)] = hist.get(str(k),0)

for sample in list(counts.keys()):

    # use sampled bit string x to compute C(x)
    x         = [int(num) for num in list(sample)]
    tmp_eng   = cost_function_C(x,G)
    
    # compute the expectation value and energy distribution
    avr_C     = avr_C    + counts[sample]*tmp_eng
    hist[str(round(tmp_eng))] = hist.get(str(round(tmp_eng)),0) + counts[sample]
    
    # save best bit string
    if( max_C[1] < tmp_eng):
        max_C[0] = sample
        max_C[1] = tmp_eng
                
M1_sampled   = avr_C/shots
""
print('\n --- SIMULATION RESULTS ---\n')
print('The sampled mean value is M1_sampled = %.02f while the true value is M1 = %.02f \n' % (M1_sampled,np.amax(F1)))
print('The approximate solution is x* = %s with C(x*) = %d \n' % (max_C[0],max_C[1]))
print('The cost function is distributed as: \n')
plot_histogram(hist,figsize = (8,6),bar_labels = False)

colors = ['r' if max_C[0][i] == '0' else 'c' for i in range(n)]
draw_graph(G, colors, pos)

# Use the IBMQ essex device
provider = IBMQ.load_account()

large_enough_devices = provider.backends(filters=lambda x: x.configuration().n_qubits >= n and
                                                   not x.configuration().simulator)
backend = least_busy(large_enough_devices)

shots = 8192

TQAOA = transpile(QAOA, backend)
qobj = assemble(TQAOA, shots=shots)
job_exp = backend.run(qobj)
job_monitor(job_exp)


exp_results = job_exp.result()
plot_histogram(exp_results.get_counts(),figsize = (10,8),bar_labels = False)


# Evaluate the data from the experiment
counts_exp = exp_results.get_counts()

avr_C_exp       = 0
max_C_exp       = [0,0]
hist_exp        = {}

for k in range(len(G.edges())+1):
    hist_exp[str(k)] = hist_exp.get(str(k),0)

for sample in list(counts_exp.keys()):

    # use sampled bit string x to compute C(x)
    x         = [int(num) for num in list(sample)]
    tmp_eng   = cost_function_C(x,G)
    
    # compute the expectation value and energy distribution
    avr_C_exp     = avr_C_exp    + counts_exp[sample]*tmp_eng
    hist_exp[str(round(tmp_eng))] = hist_exp.get(str(round(tmp_eng)),0) + counts_exp[sample]
    
    # save best bit string
    if( max_C_exp[1] < tmp_eng):
        max_C_exp[0] = sample
        max_C_exp[1] = tmp_eng
                
M1_sampled_exp   = avr_C_exp/shots

print('\n --- EXPERIMENTAL RESULTS ---\n')
print('The sampled mean value is M1_sampled = %.02f while the true value is M1 = %.02f \n' % (M1_sampled_exp,np.amax(F1)))
print('The approximate solution is x* = %s with C(x*) = %d \n' % (max_C_exp[0],max_C_exp[1]))
print('The cost function is distributed as: \n')
plot_histogram(hist_exp,figsize = (8,6),bar_labels = False)


colors = ['c' if max_C_exp[0][i] == '0' else 'r' for i in range(n)]
draw_graph(G, colors, pos)

from qiskit import Aer
from qiskit.optimization.applications.ising import max_cut
from qiskit.aqua.algorithms import QAOA as QAOAlib
from qiskit.aqua.components.optimizers import ADAM
from qiskit.aqua import aqua_globals
from qiskit.aqua import QuantumInstance
from qiskit.optimization.applications.ising.common import sample_most_likely

qubitOp, offset = max_cut.get_operator(W)

aqua_globals.random_seed = np.random.default_rng(123)
seed = 10598
backend = Aer.get_backend('statevector_simulator')
quantum_instance = QuantumInstance(backend, seed_simulator=seed, seed_transpiler=seed)

adam = ADAM(maxiter=500)
qaoa=QAOAlib(operator=qubitOp,optimizer=adam,p=1,quantum_instance=quantum_instance)
result=qaoa.run(quantum_instance)

# print results
x = sample_most_likely(result.eigenstate)
print('energy:', result.eigenvalue.real)
print('time:', result.optimizer_time)
print('max-cut objective:', result.eigenvalue.real + offset)
print('solution:', max_cut.get_graph_solution(x))
print('solution objective:', max_cut.max_cut_value(x, W))

# plot results
colors = ['r' if max_cut.get_graph_solution(x)[i] == 0 else 'c' for i in range(n)]
draw_graph(G, colors, pos)