import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from scipy.optimize import minimize
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, execute, Aer

style = {'backgroundcolor': 'lightyellow'} # Style of the drawing of the quantum circuit 

Graph_Example_1 = nx.Graph()
Graph_Example_1.add_edges_from([[1,2],[2,3],[1,3]])
nx.draw(Graph_Example_1,with_labels=True,font_weight='bold')

nx.draw(Graph_Example_1,node_color = ['b','b','r'])

Graph_Example_2 = nx.Graph()
Graph_Example_2.add_edges_from([[1,2],[1,3],[4,2],[4,3]])
nx.draw(Graph_Example_2,with_labels=True,font_weight='bold')

nx.draw(Graph_Example_2,node_color = ['r','b','b','r'])

Graph_Example_3 = nx.Graph()
Graph_Example_3.add_edges_from([[1,4],[1,2],[2,5],[4,5],[5,3],[2,3]])
nx.draw(Graph_Example_3,with_labels=True,font_weight='bold')

nx.draw(Graph_Example_3,node_color = ['r','b','b','r','b'])



Graph = nx.Graph()
Graph.add_edges_from([[0,1],[0,3],[3,4],[1,4],[1,2],[4,2]])
nx.draw(Graph,with_labels=True)

Graph.edges()

N = Graph.number_of_nodes() # The number of nodes in the graph
qc = QuantumCircuit(N,N) # Quantum circuit with N qubits and N classical register

gamma = np.pi/8 # parameter gamma for the cost function (maxcut hamiltonian)
#Apply the operator to each edge
for i, j in Graph.edges(): 
    qc.cx(i,j)
    qc.rz(2*gamma, j)
    qc.cx(i,j)
    qc.barrier()
    
qc.draw(output='mpl', style=style) # Draw the quantum circuit

beta = np.pi/8 # parameter for the mixer hamiltonian
for n in Graph.nodes():
    qc.rx(2*beta, n)

# measure the result
qc.barrier(range(N))
qc.measure(range(N), range(N))

qc.draw(output='mpl', style=style,fold= 35) 

def qaoa_circuit(G, theta,p):
    r"""
    ############################################
    # Compute the QAOA circuit for the graph G
    #
    # |psi(theta) > = |psi(beta,gamma) > = e^{-i*beta_p*B} e^{-i*gamma_p*C}... e^{-i*beta_1*B} e^{-i*gamma_1*C} H^{otimes n}|0>
    # 
    # where C is the Maxcut hamiltonian, B is the mixer hamiltonian.
    #
    # G: graph
    # theta: parameters,first half is betas, second half is gammas
    # p: number of QAOA steps
    # return: The QAOA circuit 
    ###########################################
    """
    beta = theta[:p] #Parameters beta for the mixer hamiltonian
    gamma = theta[p:]#Parameters gamma for the maxcut hamiltonian
    N = G.number_of_nodes()#Number of nodes of the graph G
    qc = QuantumCircuit(N,N)# Quantum circuit wiht N qubits and N classical register
    qc.h(range(N))# Apply the Hadamard gate to all the qubits
    # Apply p operators
    for k in range(p):
        for i, j in G.edges(): #Representation of the maxcut hamiltonian
            qc.cx(i,j)  
            qc.rz(2*gamma[k], j)
            qc.cx(i,j)
            qc.barrier()
        for n in G.nodes(): # Representation of the mixer hamiltonian
            qc.rx(2*beta[k], n)
    # Measurement in all the qubits
    qc.barrier(range(N))
    qc.measure(range(N), range(N))
    return qc #Returns the quantum circuit

# p is the number of QAOA alternating operators
p = 1

theta = np.array([np.pi/8,np.pi/8])

circuit = qaoa_circuit(Graph, theta,p)

circuit.draw(output='mpl', style=style,fold= 35) # Draw the quantum circuit

#Qiskit uses the least significant bit in the bistring, we need to invert this string
def invert_counts(counts):
    r"""
    ############################################
    # Inverted the binary string
    #
    # counts: The result of running the quantum circuit
    # return: The results with binary string inverted
    ###########################################
    """
    inv_counts = {}
    for k, v in counts.items():
        inv_counts[k[::-1]] = v
    return inv_counts 

def maxcut_objective_function(x,G):
    r"""
    ############################################
    # Compute the maxcut objective function for the binary string x and graph G
    #
    # G: graph
    # x: binary string
    # return: The cost of the binary string x 
    ###########################################
    """
    cut = 0
    for i, j in G.edges():
        if x[i] != x[j]:
            cut -= 1# the edge is cut
    return cut # Sum of all the edges that are cut

def function_to_minimize(G,p):
    r"""
    ############################################
    # Define the function to minimize given the graph G and the number of parameters
    #
    # G: graph
    # p: number of parameters
    # return: The function to minimize
    ###########################################
    """
    backend = Aer.get_backend('qasm_simulator')
    def f(theta):
        r"""
        ############################################
        # Function that gives the energy for parameters theta
        #
        # theta: parameters
        # return: Energy
        ###########################################
        """
        qc = qaoa_circuit(G,theta,p) # Gives the circuit of QAOA fot a graph G, parameters theta and p operators
        counts = execute(qc, backend, seed_simulator=10).result().get_counts()# Run the circuit and get the results using the backedn "qasm_simulator"
        
        E = 0 # Initial energy
        total_counts = 0 # Initial counts
        #Qiskit uses the least significant bit in the bistring, we need to invert this string
        InvertedCounts = invert_counts(counts) #Inverted the binary string
        # meas is the binary string measure
        # meas_count is the number of times that the binary string has been measure
        for meas, meas_count in InvertedCounts.items():
            objective_function = maxcut_objective_function(meas, G)#Compute the maxcut objetive function fot the binary string meas
            E += objective_function*meas_count # Define the energy as the value of the objective functions times the counts of that measure
            total_counts += meas_count # The total number of measures in the system
        return E/total_counts # Returns the energy given the parameters theta
    return f

p = 5 # p is the number of QAOA alternating operators
objective_function = function_to_minimize(Graph, p)
x0=np.random.randn(10)#Initial parameters

minim_solution = minimize(objective_function, x0) # Minimization of scalar function of one or more variable from scipy
minim_solution

optimal_theta = minim_solution['x'] # The optimal parameters to use

optimal_theta

backend = Aer.get_backend('qasm_simulator')
qc = qaoa_circuit(Graph, optimal_theta,p) # Define the quantum circuit with the optimal parameters
counts = invert_counts(execute(qc, backend).result().get_counts()) # Get the results
#energies = {bs: maxcut_objective_function(bs, Graph) for bs, _ in counts.items()}
#energies_sorted = {bs: en for bs, en in sorted(energies.items(), key=lambda item: item[1])}
#print(energies_sorted)
#print({bs: counts[bs] for bs, energy in energies_sorted.items()})

# Compute the maxcut objective function for the optimal parameters
results = []
for x in counts.keys():
    results.append([maxcut_objective_function(x,Graph),x]) 

# Get the best solution given the results of the maxcut objective function
best_cut, best_solution = min(results)
print(f"Best string: {best_solution} with cut: {-best_cut}")

# Define the colors of the nodes for the best solution
colors = ['r' if best_solution[node] == '0' else 'b' for node in Graph]

nx.draw(Graph,node_color = colors)



def solution_max_cut(G):
    r"""
    ############################################
    # Compute the solution for the maxcut problem given a graph G
    #
    # G: graph
    # return:The draw of the best solution with two colors. Print the solution in the binary string form and the number of cuts 
    ###############
    """
    p = 5 # p is the number of QAOA alternating operators
    objective_function = function_to_minimize(G, p)
    x0=np.random.randn(10)#Initial parameters

    new_parameters = minimize(objective_function, x0) # Minimization of scalar function of one or more variable from scipy
    
    optimal_theta = new_parameters['x']# The optimal parameters to use
    qc = qaoa_circuit(G, optimal_theta,p) # Define the quantum circuit with the optimal parameters
    counts = invert_counts(execute(qc, backend).result().get_counts()) # Get the results
    
    # Compute the maxcut objective function for the optimal parameters
    results = []
    for x in counts.keys():
        results.append([maxcut_objective_function(x,G),x])
        
    best_cut, best_solution = min(results) # Get the best solution given the results of the maxcut objective function
    colors = ['r' if best_solution[node] == '0' else 'b' for node in G] # Define the colors of the nodes for the best solution
    return nx.draw(G,node_color = colors),print(f"Best string: {best_solution} with cut: {-best_cut}")

Graph1 = nx.Graph()
Graph1.add_edges_from([[0,1],[0,2],[2,1]])
nx.draw(Graph1,with_labels=True)

solution_max_cut(Graph1)

graph2 = nx.random_regular_graph(3, 4, seed=1234)
nx.draw(graph2)          #drawing the graph
plt.show()              #plotting the graph

solution_max_cut(graph2)

graph3 = nx.random_regular_graph(3, 12, seed=1234)
nx.draw(graph3)          #drawing the graph
plt.show()              #plotting the graph

solution_max_cut(graph3)

