from just_time_windows.google_solver.google_model import evaluate_google_model
from just_time_windows.Actor.actor import Actor as NN_Actor
from just_time_windows.build_data import Raw_VRP_Data
from just_time_windows.dataloader import VRP_Dataset

def main():
    # إعداد مجموعة بيانات صغيرة للاختبار
    dataset = VRP_Dataset(dataset_size=10, num_depots=1, num_nodes=12)

    # استخراج دفعة واحدة للاختبار
    batch = dataset.get_batch(start_index=0, batch_size=10)

    # تهيئة نموذج الشبكة العصبية
    nn_actor = NN_Actor(model=None, num_movers=10, num_neighbors_action=1)

    # حساب مخرجات NN
    nn_output = nn_actor(batch)
    total_time_nn = nn_output['total_time']
    arrival_times_nn = nn_output['arrival_times']

    # استخدام Google OR-Tools لتقييم نفس البيانات
    google_output = evaluate_google_model(dataset)

    # طباعة النتائج للمقارنة
    print("Arrival times (NN):")
    print(arrival_times_nn)

    print("\nAverage Total Time (NN Actor):", total_time_nn.mean().item())
    print("Average Total Time (Google OR-Tools):", google_output.mean().item())


if __name__ == '__main__':
    main()