import os
import argparse

import json

import matplotlib
# matplotlib.use('TkAgg')
import matplotlib.pyplot as plt

from utils import visualize_tsp

from model.anneal_model import SimAnneal
from model.my_model import MyModel

def write_to_file(solution, fitness):
    with open('output/' + os.path.splitext(os.path.basename(args.tsp_file))[0] + '.txt', "w") as outfile:
        outfile.write(", ".join(str(item) for item in solution))
        outfile.write("\n")
        outfile.write(str(fitness))

def plot_learning(fitness_list):
    """
    Plot the fitness through iterations.
    """
    plt.plot([i for i in range(len(fitness_list))], fitness_list)
    plt.ylabel("Fitness")
    plt.xlabel("Iteration")
    plt.show()

def load_data(file):
    coords = []
    with open(file, "r") as infile:
        line = infile.readline()
        # Skip instance header
        while "NODE_COORD_SECTION" not in line:
            line = infile.readline()

        for line in infile.readlines():
            line = line.replace("\n", "")
            if line and 'EOF' not in line:
                line = [float(x) for x in line.lstrip().split(" ", 1)[1].split(" ") if x]
                coords.append(line)
    return coords

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='TSP Solver')
    parser.add_argument(
        'tsp_file',
        type=str,
        help='Path to tsp file.'
    )
    parser.add_argument(
        '--plot',
        action="store_true",
        help='Plot results'
    )
    args = parser.parse_args()

    coords = load_data(args.tsp_file)

    # model = MyModel(coords)
    # best_solution, fitness_list = model.fit(max_it=100000)

    sa = SimAnneal(coords)
    best_solution, fitness_list = sa.fit(max_it=100000)

    if best_solution:
        if args.plot:
            visualize_tsp.plotTSP([best_solution], coords)
            if fitness_list:
                plot_learning(fitness_list)
        print("Writing the best solution to file" )
        write_to_file(best_solution, sa.fitness(best_solution))

    data = {}
    data['nodes'] = []
    for i in range(len(coords)):
        data['nodes'].append({
            'title': str(i),
            'id': i,
            'x': int(coords[i][0]),
            'y': int(coords[i][1])
        })
    data['edges'] = []
    for i in range(len(best_solution)):
        if i == len(best_solution)-1:
            data['edges'].append({
                'source': best_solution[i],
                'target': best_solution[0]
            })
        else:
            data['edges'].append({
                'source': best_solution[i],
                'target': best_solution[i+1]
            })
    
    with open('output/' + os.path.splitext(os.path.basename(args.tsp_file))[0] + '.json', 'w') as outfile:
        json.dump(data, outfile)