import networkx as nx
from qiskit import QuantumCircuit, Aer
from qiskit.circuit import Parameter
import numpy as np
import matplotlib.pyplot as plt

from IPython.display import Image
Image('./max_cut1.png')

Image('./max_cut2.png')

Image('./mac_cut3.png')

G = nx.Graph()
G.add_nodes_from([0, 1, 2, 3])
G.add_edges_from([(0, 1), (1, 2),(2, 3), (3, 0)])
pos=nx.spring_layout(G,seed = 1)
nx.draw(G, with_labels=True, alpha=0.8, node_size=500, pos = pos)


qaoa_circuit = QuantumCircuit(len(G.nodes()))

qc = QuantumCircuit (2)
gamma = Parameter ('gamma')
qc.rzz(2*gamma, 0, 1)
qc.decompose().draw()

def initial_state(circuit):
    for i in range(circuit.num_qubits):
        circuit.h(i)
    return circuit

qaoa_circuit = initial_state(qaoa_circuit)
qaoa_circuit.draw()

def problem_unitary(graph,circuit,gamma):
    for edge in list(G.edges()):
        circuit.rzz(2*gamma,edge[0],edge[1])
    return circuit
qaoa_circuit.barrier()
qaoa_circuit = problem_unitary(G, qaoa_circuit, gamma)
qaoa_circuit.decompose().draw()

beta = Parameter('beta')
def mix_unitary(circuit,beta):
    for i in range(circuit.num_qubits):
        circuit.rx(2*beta,i)
    return circuit
qaoa_circuit.barrier()
qaoa_circuit = mix_unitary(qaoa_circuit, beta)
qaoa_circuit.draw()

def create_qaoa_circuit(graph, theta):
    p = len(theta)//2
    gamma = theta[p:]
    beta = theta[:p]
    qaoa_p = QuantumCircuit(len(graph.nodes()))
    qaoa_p = initial_state(qaoa_p)
    for i in range(p):
        qaoa_p = problem_unitary(graph, qaoa_p, gamma[i])
        qaoa_p = mix_unitary(qaoa_p, beta[i])
    qaoa_p.measure_all()
    return qaoa_p
    
    

def get_cost(graph, state):
    cost = 0
    for i,j in graph.edges():
        if state[i] != state[j]:
            cost -= 1
    
    return cost

def expectation_value(graph, counts):
    avg = 0
    sum_count = 0
    for state, count in counts.items():
        cost = get_cost(graph, state)
        avg += cost*count
        sum_count += count
        
    return avg/sum_count

def run_qpu(graph, shots=512):
    backend = Aer.get_backend('aer_simulator')
    backend.shots = shots
    
    def run_circuit(theta):
        qc = create_qaoa_circuit(graph, theta)
        counts = backend.run(qc, seed_simulator = 5, nshots= shots).result().get_counts()
        return expectation_value(graph,counts)
    
    return run_circuit

from scipy.optimize import minimize

p=1
theta = [1.0 for i in range(2*p)]
exp_value = run_qpu(G)

res = minimize(exp_value, theta , method = 'COBYLA')
res

from qiskit.visualization import plot_histogram

backend = Aer.get_backend('aer_simulator')
backend.shots = 512

qc_res = create_qaoa_circuit(G, res.x)

counts = backend.run(qc_res, seed_simulator = 5).result().get_counts()

plot_histogram(counts)