from __future__ import print_function
from ortools.constraint_solver import routing_enums_pb2
from ortools.constraint_solver import pywrapcp
from just_time_windows.dataloader import VRP_Dataset
from just_time_windows.google_solver.convert_data import convert_data


def compute_total_time(data, routing, assignment):
    time_dimension = routing.GetDimensionOrDie('Time')
    total_time = 0
    for vehicle_id in range(data['num_vehicles']):
        index = routing.Start(vehicle_id)
        while not routing.IsEnd(index):
            index = assignment.Value(routing.NextVar(index))
        time_var = time_dimension.CumulVar(index)
        total_time += assignment.Min(time_var)
    return total_time


def print_solution(data, manager, routing, assignment):
    time_dimension = routing.GetDimensionOrDie('Time')
    total_time = 0
    for vehicle_id in range(data['num_vehicles']):
        index = routing.Start(vehicle_id)
        plan_output = f'Route for vehicle {vehicle_id}:\n'
        while not routing.IsEnd(index):
            time_var = time_dimension.CumulVar(index)
            plan_output += f'{manager.IndexToNode(index)} Time({assignment.Min(time_var)},{assignment.Max(time_var)}) -> '
            index = assignment.Value(routing.NextVar(index))
        time_var = time_dimension.CumulVar(index)
        plan_output += f'{manager.IndexToNode(index)} Time({assignment.Min(time_var)},{assignment.Max(time_var)})\n'
        plan_output += f'Time of the route: {assignment.Min(time_var)} min\n'
        print(plan_output)
        total_time += assignment.Min(time_var)
    print(f'Total time of all routes: {total_time} min')


def adjust_time_windows(windows, offset=10):
    return [(start + offset, end + offset) for start, end in windows]


def main():
    # إعداد مجموعة بيانات لاختبار نموذج time windows
    dataset = VRP_Dataset(dataset_size=1, num_nodes=20, num_depots=1)
    batch = dataset.get_batch(0, 1)
    data = convert_data(batch, scale_factor=100)[0]

    manager = pywrapcp.RoutingIndexManager(
        len(data['time_matrix']),
        data['num_vehicles'],
        data['depot']
    )
    routing = pywrapcp.RoutingModel(manager)

    def time_callback(from_index, to_index):
        from_node = manager.IndexToNode(from_index)
        to_node = manager.IndexToNode(to_index)
        return data['time_matrix'][from_node][to_node]

    transit_callback_index = routing.RegisterTransitCallback(time_callback)
    routing.SetArcCostEvaluatorOfAllVehicles(transit_callback_index)

    routing.AddDimension(
        transit_callback_index,
        10000,
        10000,
        False,
        'Time'
    )

    time_dimension = routing.GetDimensionOrDie('Time')

    # تعيين نوافذ زمنية للعقد
    for location_idx, time_window in enumerate(data['time_windows']):
        index = manager.NodeToIndex(location_idx)
        time_dimension.CumulVar(index).SetRange(int(time_window[0]), int(time_window[1]))

    # تعيين نوافذ زمنية لنقاط بداية المركبات
    for vehicle_id in range(data['num_vehicles']):
        index = routing.Start(vehicle_id)
        depot_window = data['time_windows'][0]
        time_dimension.CumulVar(index).SetRange(int(depot_window[0]), int(depot_window[1]))

    for i in range(data['num_vehicles']):
        routing.AddVariableMinimizedByFinalizer(time_dimension.CumulVar(routing.Start(i)))
        routing.AddVariableMinimizedByFinalizer(time_dimension.CumulVar(routing.End(i)))

    search_parameters = pywrapcp.DefaultRoutingSearchParameters()
    search_parameters.first_solution_strategy = routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC

    assignment = routing.SolveWithParameters(search_parameters)

    if assignment:
        print_solution(data, manager, routing, assignment)
        total_time = compute_total_time(data, routing, assignment)
        print(f"\nFinal total time: {total_time} min")
    else:
        print("No solution found.")


if __name__ == '__main__':
    main()