import os
import signal
import json
import math
import timeit

from utils.load_data import load_data
from utils.load_data import log

def timeout_handler(signum, frame):
    raise Exception("Timeout")

def dist(node_0, node_1, coords):
        """
        Euclidean distance between two nodes.
        """
        coord_0, coord_1 = coords[node_0], coords[node_1]
        return math.sqrt((coord_0[0] - coord_1[0]) ** 2 + (coord_0[1] - coord_1[1]) ** 2)

def fitness(solution, coords):
    N = len(coords)
    if(len(solution) != N):
        return math.inf
    cur_fit = 0
    for i in range(len(solution)):
        cur_fit += dist(solution[i % N], solution[(i + 1) % N], coords)
    return cur_fit

def TSP(tsp_file, model, timeout=60):

    best_solution = []
    fitness_list = []

    nodes = load_data(tsp_file)

    model = model()
    # Set timeout
    signal.signal(signal.SIGALRM, timeout_handler)
    signal.alarm(timeout)

    if not os.path.exists('output'):
        os.makedirs('output')

    # Try your algorithm
    try:
        model.init(nodes)
        best_solution, fitness_list = model.fit()
    except Exception as exc: 
        if(str(exc) == "Timeout"):
            log("Timeout -3")
            with open('output/' + os.path.splitext(os.path.basename(tsp_file))[0] + '.txt', "w") as outfile:
                outfile.write("-3")
        elif (len(best_solution) == 0):
            print(exec)
            log("No Answer -1")
            with open('output/' + os.path.splitext(os.path.basename(tsp_file))[0] + '.txt', "w") as outfile:
                outfile.write("-1")
        else:
            print(exc)
            log("Unkown -4")
            with open('output/' + os.path.splitext(os.path.basename(tsp_file))[0] + '.txt', "w") as outfile:
                outfile.write("-4")

        best_solution = [-1] * len(nodes)
        return best_solution, fitness_list
    
    signal.alarm(0)

    # Collect results
    if(len(fitness_list) > 0):
        log("[Node] " + str(len(best_solution)) + ", [Best] " + str(fitness_list[fitness_list.index(min(fitness_list))]))

    if (len(best_solution) != len(nodes)):
        print(len(best_solution))
        log("Invalid -2")
        with open('output/' + os.path.splitext(os.path.basename(tsp_file))[0] + '.txt', "w") as outfile:
            outfile.write("-2")
        best_solution = [-1] * len(nodes)
    else:
        # log("Writing the best solution to file" )
        with open('output/' + os.path.splitext(os.path.basename(tsp_file))[0] + '.txt', "w") as outfile:
            outfile.write(str(model.fitness(best_solution)))
            outfile.write("\n")
            outfile.write(", ".join(str(item) for item in best_solution))

        # Write to JSON
        data = {}

        data['nodes'] = []
        for i in range(len(nodes)):
            data['nodes'].append({
                'title': str(i),
                'id': i,
                'x': int(nodes[i][0]),
                'y': int(nodes[i][1])
            })

        data['edges'] = []
        for i in range(len(best_solution)):
            data['edges'].append({
                'source': best_solution[i % len(best_solution)],
                'target': best_solution[(i+1) % len(best_solution)]
            })

        with open('output/' + os.path.splitext(os.path.basename(tsp_file))[0] + '.json', 'w') as outfile:
            json.dump(data, outfile)

    return best_solution, fitness_list

def TSP_Bench_ONE(tsp_file, model, timeout=60):

    start = timeit.default_timer()

    # Model Running
    best_solution, fitness_list = TSP(tsp_file, model, timeout=timeout)
    # Model End

    stop = timeit.default_timer()
    print('[*] Running for: {time:.2f} seconds'.format(time=(stop - start)))

    print()
    return best_solution, fitness_list, (stop - start)

def TSP_Bench_PATH(tsp_file_path, model, timeout=60):
    best_solutions = []
    fitness_lists = []
    times = []
    for root, _, files in os.walk(tsp_file_path):
        if(files):
            for f in files:
                # Get input file name
                tsp_file = os.path.join(root, str(f))
                log(tsp_file)

                # Run TSP
                best_solution, fitness_list, time = TSP_Bench_ONE(tsp_file, model, timeout=timeout)
                best_solutions.append(best_solution)
                fitness_lists.append(fitness_lists)
                times.append(time)

    return best_solutions, fitness_lists, times

def TSP_Bench_ALL(root, model):
    best_solutions_all = []
    fitness_lists_all = []
    times_all = []

    best_solutions_simple, fitness_lists_simple, times_simple = TSP_Bench_PATH(os.path.join(root, "data/simple/"), model, timeout=60)
    best_solutions_medium, fitness_lists_medium, times_medium = TSP_Bench_PATH(os.path.join(root, "data/medium/"), model, timeout=180)
    best_solutions_hard, fitness_lists_hard, times_hard = TSP_Bench_PATH(os.path.join(root, "data/hard/"), model, timeout=300)

    best_solutions_all.append(best_solutions_simple)
    best_solutions_all.append(best_solutions_medium)
    best_solutions_all.append(best_solutions_hard)

    fitness_lists_all.append(fitness_lists_simple)
    fitness_lists_all.append(fitness_lists_medium)
    fitness_lists_all.append(fitness_lists_hard)

    times_all.append(times_simple)
    times_all.append(times_medium)
    times_all.append(times_hard)

    return [item for sublist in best_solutions_all for item in sublist], [item for sublist in fitness_lists_all for item in sublist], [item for sublist in times_all for item in sublist]