import argparse
import pprint as pp
import time
import warnings
from multiprocessing import Pool

import lkh
import scipy.spatial
import numpy as np
import tqdm
import tsplib95

class TSPEvaluator(object):
  def __init__(self, points):
    self.dist_mat = scipy.spatial.distance_matrix(points, points)

  def evaluate(self, route):
    total_cost = 0
    for i in range(len(route) - 1):
      total_cost += self.dist_mat[route[i]-1, route[i + 1]-1]
    return total_cost

num_nodes = 5000
lkh_trails = 10000
filename = f"./tsp/tsp{num_nodes}_test_lkh.txt"
# filename = "./tsp/tsp50_test_concorde.txt"

with open(filename, "r") as f:
    filelines = f.readlines()

res = []

t1 = time.time()
for i in tqdm.tqdm(range(len(filelines))):
# for i in range(len(filelines)):
    line = filelines[i]
    line = line.strip()
    points = line.split(' output ')[0]
    points = points.split(' ')
    # points = points.split()
    nodes_coord = np.array([[float(points[i]), float(points[i + 1])] for i in range(0, len(points), 2)])
    scale = 1e6
    lkh_path = './LKH-3.0.6/LKH'
    problem = tsplib95.models.StandardProblem()
    problem.name = 'TSP'
    problem.type = 'TSP'
    problem.dimension = num_nodes
    problem.edge_weight_type = 'EUC_2D'
    problem.node_coords = {n + 1: nodes_coord[n] * scale for n in range(num_nodes)}

    solution = lkh.solve(lkh_path, problem=problem, max_trials=lkh_trails, runs=1)
    tsp_eval = TSPEvaluator(nodes_coord)
    optimal_val = tsp_eval.evaluate(solution[0])
    print("The solution is: ")
    print(optimal_val)
    res.append(optimal_val)
    # exit(0)

t2 = time.time()
print(sum(res) / len(res))
print(t2 - t1)
