import qiskit
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from collections import defaultdict
from operator import itemgetter
from scipy.optimize import minimize
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, execute, Aer
from qiskit.aqua.algorithms import NumPyEigensolver
from qiskit.quantum_info import Pauli
from qiskit.aqua.operators import op_converter
from qiskit.aqua.operators import WeightedPauliOperator

from tsp_qaoa import marina_solution



G=nx.Graph()
i=1
G.add_node(i,pos=(i,i))
G.add_node(2,pos=(2,2))
G.add_node(3,pos=(1,0))
G.add_edge(1,2,weight=20.5)
G.add_edge(1,3,weight=9.8)
pos=nx.get_node_attributes(G,'pos')
nx.draw(G,pos)
labels = nx.get_edge_attributes(G,'weight')
nx.draw_networkx_edge_labels(G,pos,edge_labels=labels)

def append_zz_term(qc,q1,q2,gamma):
    qc.cx(q1,q2)
    qc.rz(2*gamma,q2)
    qc.cx(q1,q2)


def get_cost_circuit(G,gamma):
    N=G.number_of_nodes()
    qc=QuantumCircuit(N,N)
    for i,j in G.edges():
        append_zz_term(qc,i,j,gamma)
    return qc
#print(get_cost_circuit(G,0.5))

def append_x_term(qc,q1,beta):
    qc.rx(2*beta,q1)

def get_mixer_operator(G,beta):
    N=G.number_of_nodes()
    qc=QuantumCircuit(N,N)
    for n in G.nodes():
        append_x_term(qc,n,beta)
    return qc

#print(get_mixer_operator(G,0.5))


def get_QAOA_circuit(G,beta,gamma):
    assert(len(beta)==len(gamma))
    N=G.number_of_nodes()
    qc=QuantumCircuit(N,N)
    qc.h(range(N))
    p=len(beta)
    #aplicamos las p rotaciones
    for i in range(p):
        qc=qc.compose(get_cost_circuit(G,gamma[i]))
        qc=qc.compose(get_mixer_operator(G,beta[i]))
    qc.barrier(range(N))
    qc.measure(range(N),range(N))
    return qc
    
print(get_QAOA_circuit(G,[0.5,0,6],[0.5,0,6]))    

        

def invert_counts(counts):
    return {k[::-1] :v for k,v in counts.items()}

qc=get_QAOA_circuit(G,[0.5,0,6],[0.5,0,6])   
backend=Aer.get_backend('qasm_simulator')
job=execute(qc,backend)
result=job.result()
print(invert_counts(result.get_counts()))

def maxcut_obj(x,G):
    cut=0
    for i,j in G.edges():
        if x[i]!=x[j]:
            cut = cut-1
    return cut
print(maxcut_obj("00011",G))

def compute_maxcut_energy(counts,G):
    energy=0
    get_counts=0
    total_counts=0
    for meas, meas_count in counts.items():
        obj_for_meas=maxcut_obj(meas,G)
        energy+=obj_for_meas*meas_count
        total_counts+=meas_count
    return energy/total_counts

def get_black_box_objective(G,p):
    backend=Aer.get_backend('qasm_simulator')
    def f(theta):
        beta=theta[:p]
        gamma=theta[p:]
        qc=get_QAOA_circuit(G,beta,gamma)
        counts=execute(qc,backend,seed_simulator=10).result().get_counts()
        return compute_maxcut_energy(invert_counts(counts),G)
    return f
p=5
obj=get_black_box_objective(G,p)
init_point=np.array([0.8,2.2,0.83,2.15,0.37,2.4,6.1,2.2,3.8,6.1])#([2,2,1,1,1,1,1,1,1,1])
res_sample=minimize(obj, init_point,method="COBYLA",options={"maxiter":2500,"disp":True})

res_sample
        

from thirdParty.classical import rand_graph, classical, bitstring_to_path, calc_cost
from utils import mapeo_grafo

cantidad_ciudades = 4
pesos, conexiones = None, None
mejor_camino = None
while not mejor_camino:
    pesos, conexiones = rand_graph(cantidad_ciudades)
    mejor_costo, mejor_camino = classical(pesos, conexiones, loop=False)
        

G = mapeo_grafo(conexiones, pesos)

pos=nx.spring_layout(G)
nx.draw(G,pos)
labels = nx.get_edge_attributes(G,'weight')
nx.draw_networkx_edge_labels(G,pos,edge_labels=labels)


G

pos=nx.get_node_attributes(G,'weight')

pos

labels = nx.get_edge_attributes(G,'weight')

labels

def funcion_costo(multiplicador_lagrange, cantidad_ciudades, pesos, conexiones ):
    N = G.number_of_nodes()
    N_square = N^2
    # restriccion 1
    for i in range(cantidad_ciudades):
        cur = sI(N_square) 
        for j in range(num_cities):
            cur -= D(i, j)  
        ret += cur**2
        
    

# retorna el indice de qubit por conversion al problema
def quibit_indice(i, l, N):
    return i * N + l



    
















from qiskit.quantum_info.operators import Operator, Pauli
# Create an operator
XX = Operator(Pauli(label='XX'))

# Add to a circuit
circ = QuantumCircuit(2, 2)
circ.append(XX, [0, 1])
circ.measure([0,1], [0,1])
circ.draw('mpl')




# Add to a circuit
circ = QuantumCircuit(2, 2)
circ.append(a, [0])
#circ.measure([0,1], [0,1])
circ.draw('mpl')

a = I - ( 0.5*(I+ Z))**2
a = Operator(a)

a.is_unitary()

print(I @ Z)

