import sys
import time
import torch
import argparse
import utils
import multiprocessing as mp
from concurrent.futures import ThreadPoolExecutor

from ortools.constraint_solver import pywrapcp
from ortools.constraint_solver import routing_enums_pb2


def solve(problem, i, max_time):
    scale = 100000
    size = problem.task_demand.size(1)
    demand = [0] + problem.task_demand[i].tolist()
    capacity = problem.worker_weight_limit[i].tolist()
    distance = (problem.distance_matrix[i] * scale + 0.5).to(torch.int32).tolist()

    queue = mp.Queue()
    p = mp.Process(target=do_solve, args=(size, demand, capacity, distance, max_time, queue))
    p.start()
    p.join()

    return queue.get() / scale, queue.get()


def do_solve(size, demand, capacity, distance, max_time, queue):
    capacity = capacity * size

    manager = pywrapcp.RoutingIndexManager(size + 1, size, 0)
    routing = pywrapcp.RoutingModel(manager)

    def distance_callback(from_index, to_index):
        from_node = manager.IndexToNode(from_index)
        to_node = manager.IndexToNode(to_index)
        return distance[from_node][to_node]

    distance_callback_index = routing.RegisterTransitCallback(distance_callback)
    routing.SetArcCostEvaluatorOfAllVehicles(distance_callback_index)

    def demand_callback(from_index):
        from_node = manager.IndexToNode(from_index)
        return demand[from_node]

    demand_callback_index = routing.RegisterUnaryTransitCallback(demand_callback)
    routing.AddDimensionWithVehicleCapacity(demand_callback_index, 0, capacity, True, 'capacity')

    params = pywrapcp.DefaultRoutingSearchParameters()
    params.first_solution_strategy = (routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC)
    params.local_search_metaheuristic = (routing_enums_pb2.LocalSearchMetaheuristic.GUIDED_LOCAL_SEARCH)
    params.time_limit.seconds = max_time

    start_time = time.time()
    solution = routing.SolveWithParameters(params)
    spent_time = time.time() - start_time

    queue.put(solution.ObjectiveValue())
    queue.put(spent_time)


def run_orts(task, max_time):
    problem, i = task
    return solve(problem, i, max_time)


def main(args):
    print("args: {}".format(vars(args)))
    problem_size = args.problem_size
    problem_count = args.problem_count
    batch_size = args.batch_size

    assert problem_count % batch_size == 0
    batch_count = problem_count // batch_size
    problem_list = utils.make_problem(batch_count, batch_size, problem_size)

    executor = ThreadPoolExecutor(max_workers=args.threads)
    task_list = [(p, i) for p in problem_list for i in range(batch_size)]

    total_cost = 0
    total_time = 0
    for cost, elapse in executor.map(run_orts, task_list, [args.max_time] * problem_count):
        total_cost += cost
        total_time += elapse

    avg_cost = total_cost / problem_count
    avg_time = total_time / problem_count
    print()
    print("-----------------------------------------------------")
    print("avg_cost: {:.4f}".format(avg_cost))
    print("avg_time: {:.6f}s".format(avg_time))
    print("total_count: {}".format(problem_count))
    print("-----------------------------------------------------\n")
    sys.stdout.flush()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--threads', default=20, type=int, help='number of threads')
    parser.add_argument('--max_time', default=60, type=int, help='the time limit for the search in seconds')

    parser.add_argument('--problem_size', default=100, type=int, choices=[100, 1000, 2000, 5000],  help='problem size')
    parser.add_argument('--problem_count', default=128, type=int,  help='total number of generated problem instances')
    parser.add_argument('--batch_size', default=128, type=int,  help='batch size for feedforwarding')

    args = parser.parse_args()
    main(args)