import networkx as nx
import numpy as np

from qiskit import Aer
from qiskit.aqua import aqua_globals, QuantumInstance
from qiskit.aqua.algorithms import QAOA, VQE, NumPyMinimumEigensolver
from qiskit.aqua.components.optimizers import SPSA, COBYLA
from qiskit.optimization.applications.ising.common import sample_most_likely
from qiskit.optimization.applications.ising import tsp
from qiskit.optimization.converters import IsingToQuadraticProgram
from qiskit.optimization.problems import QuadraticProgram
from qiskit.optimization.algorithms import MinimumEigenOptimizer, RecursiveMinimumEigenOptimizer

from qiskit.circuit.library import TwoLocal


def get_graph(nodes):

    G = nx.Graph()
    G.add_nodes_from(np.arange(0, nodes, 1))

    # Create random positions in the graph. Distance will be calculated from positions
    # Note: Dwave and other solvers require a complete graph
    for i in range(nodes):
        G.nodes[i]['pos'] = (np.random.uniform(0, 10), np.random.uniform(0, 10))

    elist = set()
    for i in range(nodes):
        for t in range(i + 1,nodes):
            y1=G.nodes[i]['pos'][1]
            x1=G.nodes[i]['pos'][0]
            y2=G.nodes[t]['pos'][1]
            x2=G.nodes[t]['pos'][0]
            dist = np.sqrt(((x2-x1)**2)+((y2-y1)**2))
            _tuple = (i, t, dist)
            elist.add(_tuple)

    # tuple is (i,j,weight) where (i,j) is the edge
    G.add_weighted_edges_from(elist)

    return G

def get_cost_matrix(G, nodes):
    w = np.zeros([nodes,nodes])
    for i in range(nodes):
        for j in range(nodes):
            temp = G.get_edge_data(i,j,default=0)
            if temp != 0:
                w[i,j] = temp['weight']
    return w

def calculate_cost(cost_matrix, solution):
    cost = 0
    for i in range(len(solution)):
        a = i % len(solution)
        b = (i + 1) % len(solution)
        cost += cost_matrix[solution[a]][solution[b]]

    return cost



nodes = 3
starting_node = 0

####
G = get_graph(nodes)
cost_matrix = get_cost_matrix(G, nodes)

G.nodes[0]['pos']
coords = []
for node in G.nodes:
    coords.append(G.nodes[node]['pos'])
    print(G.nodes[node]['pos'])


tsp_instance = tsp.TspData(name = "TSP", dim = len(G.nodes), coord = coords, w = cost_matrix)
qubitOp, offset = tsp.get_operator(tsp_instance)
#print(qubitOp.print_details())
print("Qubits: ",qubitOp.num_qubits)

qp = QuadraticProgram()
qp.from_ising(qubitOp, offset, linear=True)
qp.to_docplex().prettyprint()

backend = Aer.get_backend('qasm_simulator')
quantum_instance = QuantumInstance(backend)

optimizer = SPSA(maxiter=400)
#optimizer = COBYLA(maxiter=100, rhobeg=2, tol=0.5, disp=True)
ry = TwoLocal(qubitOp.num_qubits, 'ry', 'cx', reps=3, entanglement='full')
vqe = VQE(operator=qubitOp, optimizer=optimizer, quantum_instance=quantum_instance)
result = vqe.run(quantum_instance)
print(result)

x = sample_most_likely(result.eigenstate)
print(x)
print(tsp.tsp_feasible(x))
if(tsp.tsp_feasible(x)):
    z = tsp.get_tsp_solution(x)
    print('solution with VQE:', z)
else:
    print('No solution found with VQE and such parameters')

npes = NumPyMinimumEigensolver(qubitOp)
result_classical = npes.run()
x = sample_most_likely(result_classical.eigenstate)
print(x)
if(tsp.tsp_feasible(x)):
    z = tsp.get_tsp_solution(x)
    print('solution Numpy classical solver:', z)
else:
    print('No solution found with Classical Eigensolver')

## Try now with the optimizer and the Recursive Optimizer

backend = quantum_instance = Aer.get_backend('qasm_simulator')
optimizer = COBYLA(maxiter=200, rhobeg=3, tol=1.5, disp=True)

# MinimumEigenSolvers
qaoa_mes = QAOA(quantum_instance = backend, optimizer=optimizer)
vqe_mes = VQE(quantum_instance=backend, optimizer=optimizer, operator=qubitOp)
exact_mes = NumPyMinimumEigensolver()

# Optimizers
qaoa = MinimumEigenOptimizer(qaoa_mes)   # using QAOA
vqe = MinimumEigenOptimizer(vqe_mes)   # using QAOA
exact = MinimumEigenOptimizer(exact_mes)  # using the exact classical numpy minimum eigen solver

# We can try using VQE or QAOA for the minimum eigen optimizer and compare results
rqaoa = RecursiveMinimumEigenOptimizer(min_eigen_optimizer=vqe, min_num_vars=1, min_num_vars_optimizer=exact)
#rqaoa = RecursiveMinimumEigenOptimizer(min_eigen_optimizer=qaoa, min_num_vars=1, min_num_vars_optimizer=exact)


exact_result = exact.solve(qp)
print("Result Exact QP: ",exact_result.x)
if(tsp.tsp_feasible(exact_result.x)):
    print(tsp.get_tsp_solution(exact_result.x))
      
qaoa_result = qaoa.solve(qp)
print("Result QAOA with QP: ",qaoa_result.x)
if(tsp.tsp_feasible(qaoa_result.x)):
    print(tsp.get_tsp_solution(qaoa_result.x))

vqe_result = vqe.solve(qp)
print("Result VQE with QP: ",vqe_result.x)
if(tsp.tsp_feasible(vqe_result.x)):
    print(tsp.get_tsp_solution(vqe_result.x))
      
rqaoa_result = rqaoa.solve(qp)
print("Result RQAOA (withVQE) QP: ",rqaoa_result.x)
if(tsp.tsp_feasible(rqaoa_result.x)):
    print(tsp.get_tsp_solution(rqaoa_result.x))

# Compare with using VQE directly without optimizer. That is why we send qubit operator to the MES
result = vqe_mes.run(quantum_instance)
x = sample_most_likely(result_classical.eigenstate)
print("Result VQE with Operator: ",x)
if(tsp.tsp_feasible(x)):
    print(tsp.get_tsp_solution(x))


