import matplotlib.pyplot as ply
%matplotlib inline

import networkx as nx
import numpy as np

from qiskit_aqua.translators.ising import tsp
from qiskit_aqua.input import EnergyInput
from qiskit_aqua import run_algorithm
from qiskit_qcgpu_provider import QCGPUProvider

locations = 3

problem = tsp.random_tsp(locations)
positions = {k: v for k, v in enumerate(problem.coord)}

G = nx.Graph()
G.add_nodes_from(np.arange(0, locations, 1))
nx.draw(G, with_labels=True, pos=positions)

best_distance, best_order = brute_force(problem.w, problem.dim)
draw(G, best_order, positions)

operator, offset = tsp.get_tsp_qubitops(problem)
algorithm_input = EnergyInput(operator)

algorithm_parameters = {
    'problem': { 'name': 'ising', 'random_seed': 23 },
    'algorithm': { 'name': 'VQE', 'operator_mode': 'matrix' },
    'optimizer': { 'name': 'SPSA', 'max_trials':100 },
    'variational_form': {'name': 'RY', 'depth': 5, 'entanglement': 'linear'}
}

backend = QCGPUProvider().get_backend('statevector_simulator')
%time result_qiskit = run_algorithm(algorithm_parameters, algorithm_input)
%time result = run_algorithm(algorithm_parameters, algorithm_input, backend=backend)


#print('tsp objective:', result['energy'] + offset)
x = tsp.sample_most_likely(result['eigvecs'][0])
print('feasible:', tsp.tsp_feasible(x))
z = tsp.get_tsp_solution(x)
print('solution:', z)
print('solution objective:', tsp.tsp_value(z, problem.w))
draw(G, z, positions)

# Utitlity Functions
def draw(G, order, positions):
    G2 = G.copy()
    n = len(order)
    for i in range(n):
        j = (i + 1) % n
        G2.add_edge(order[i], order[j])
    nx.draw(G2, pos=positions, with_labels=True)
    
# Classically solve the problem using a brute-force method
from itertools import permutations

def brute_force(weights, N):
    a = list(permutations(range(1, N)))
    best_distance = None
    for i in a:
        distance = 0
        pre_j = 0
        for j in i:
            distance += weights[j, pre_j]
            pre_j = j
        distance += weights[pre_j, 0]
        order = (0,) + i
        if best_distance is None or distance < best_distance:
            best_order = order
            best_distance = distance
            
    return best_distance, best_order

import warnings
warnings.filterwarnings('ignore')

