From c417e071bd55737c6599f2ae20107e428209e0ef Mon Sep 17 00:00:00 2001
From: Ondrej Sykora <ondrasej@google.com>
Date: Wed, 25 Oct 2023 20:40:31 +0000
Subject: [PATCH] Fixed a bug in the computation of time windows for delivery
 rounds in the local plan.

---
 examples/two_step_routing/two_step_routing.py |  17 ++-
 .../two_step_routing/two_step_routing_test.py | 129 ++++++++++++++++++
 2 files changed, 144 insertions(+), 2 deletions(-)

diff --git a/examples/two_step_routing/two_step_routing.py b/examples/two_step_routing/two_step_routing.py
index e5a11bc5..60dbe71c 100644
--- a/examples/two_step_routing/two_step_routing.py
+++ b/examples/two_step_routing/two_step_routing.py
@@ -1428,6 +1428,7 @@ def _get_local_model_route_start_time_windows(
   global_end_time = cfr_json.get_global_end_time(model)
 
   route_start_time = cfr_json.parse_time_string(route["vehicleStartTime"])
+  shipments = cfr_json.get_shipments(model)
 
   # The start time window for the route is computed as the intersection of
   # "route start time windows" of all visits in the route. A "route start time
@@ -1441,8 +1442,20 @@ def _get_local_model_route_start_time_windows(
   overall_route_start_time_intervals = ((global_start_time, global_end_time),)
 
   for visit in visits:
-    visit_request = cfr_json.get_visit_request(model, visit)
-    time_windows = visit_request.get("timeWindows")
+    # NOTE(ondrasej): We can't use `visit["shipmentIndex"]` to get the shipment;
+    # `visit` is from the local model, while `model` is the global model. To
+    # get the expected results, we need to use the shipment label from the visit
+    # to get the shipment index in the base model.
+    shipment_label = visit.get("shipmentLabel")
+    shipment_index = _get_shipment_index_from_local_label(shipment_label)
+    shipment = shipments[shipment_index]
+    deliveries = shipment.get("deliveries", ())
+    if len(deliveries) != 1:
+      raise ValueError(
+          "Only shipments with one delivery request are supported."
+      )
+
+    time_windows = deliveries[0].get("timeWindows")
     if not time_windows:
       # This shipment can be delivered at any time. No refinement of the route
       # delivery time interval is needed.
diff --git a/examples/two_step_routing/two_step_routing_test.py b/examples/two_step_routing/two_step_routing_test.py
index 176bddea..60a1328b 100644
--- a/examples/two_step_routing/two_step_routing_test.py
+++ b/examples/two_step_routing/two_step_routing_test.py
@@ -3286,6 +3286,135 @@ def test_local_refinement_model(self):
     )
 
 
+class GetLocalModelRouteStartTimeWindowsTest(unittest.TestCase):
+  """Tests for _get_local_model_route_start_time_windows."""
+
+  maxDiff = None
+
+  _MODEL: cfr_json.ShipmentModel = {
+      "globalStartTime": "2023-10-25T00:00:00Z",
+      "globalEndTime": "2023-10-25T23:59:59Z",
+      "shipments": [
+          {
+              "deliveries": [
+                  {
+                      "timeWindows": [{
+                          "startTime": "2023-10-25T09:00:00Z",
+                          "endTime": "2023-10-25T12:00:00Z",
+                      }]
+                  }
+              ],
+              "label": "S001",
+          },
+          {
+              "deliveries": [
+                  {
+                      "timeWindows": [{
+                          "startTime": "2023-10-25T09:00:00Z",
+                          "endTime": "2023-10-25T12:00:00Z",
+                      }]
+                  }
+              ],
+              "label": "S002",
+          },
+          {
+              "deliveries": [
+                  {
+                      "timeWindows": [{
+                          "startTime": "2023-10-25T14:00:00Z",
+                          "endTime": "2023-10-25T16:00:00Z",
+                      }]
+                  }
+              ],
+              "label": "S003",
+          },
+          {
+              "deliveries": [
+                  {
+                      "timeWindows": [{
+                          "startTime": "2023-10-25T12:00:00Z",
+                          "endTime": "2023-10-25T15:00:00Z",
+                      }]
+                  }
+              ],
+              "label": "S004",
+          },
+          {
+              "deliveries": [{}],
+              "label": "S005",
+          },
+          {
+              "deliveries": [{}],
+              "label": "S006",
+          },
+      ],
+  }
+
+  def test_empty_route(self):
+    self.assertIsNone(
+        two_step_routing._get_local_model_route_start_time_windows({}, {})
+    )
+
+  def test_with_invalid_route(self):
+    local_route: cfr_json.ShipmentRoute = {
+        "vehicleStartTime": "2023-10-25T11:00:00Z",
+        "vehicleLabel": "P001 []",
+        "visits": [
+            {
+                "startTime": "2023-10-25T11:10:00Z",
+                "shipmentIndex": 3,
+                "shipmentLabel": "0: S001",
+            },
+            {
+                "startTime": "2023-10-25T11:20:00Z",
+                "shipmentIndex": 1,
+                "shipmentLabel": "2: S004",
+            },
+        ],
+    }
+    with self.assertRaisesRegex(ValueError, "incompatible time windows"):
+      two_step_routing._get_local_model_route_start_time_windows(
+          self._MODEL, local_route
+      )
+
+  def test_with_some_shipments(self):
+    local_route: cfr_json.ShipmentRoute = {
+        "vehicleStartTime": "2023-10-25T11:00:00Z",
+        "vehicleLabel": "P001 []",
+        "visits": [
+            {
+                "startTime": "2023-10-25T11:10:00Z",
+                "shipmentIndex": 3,
+                "shipmentLabel": "0: S001",
+            },
+            {
+                "startTime": "2023-10-25T11:20:00Z",
+                "shipmentIndex": 1,
+                "shipmentLabel": "4: S005",
+            },
+            {
+                "startTime": "2023-10-25T11:45:00Z",
+                "shipmentIndex": 2,
+                "shipmentLabel": "1: S002",
+            },
+            {
+                "startTime": "2023-10-25T12:50:00Z",
+                "shipmentIndex": 0,
+                "shipmentLabel": "3: S004",
+            },
+        ],
+    }
+    self.assertSequenceEqual(
+        two_step_routing._get_local_model_route_start_time_windows(
+            self._MODEL, local_route
+        ),
+        [{
+            "startTime": "2023-10-25T10:10:00Z",
+            "endTime": "2023-10-25T11:15:00Z",
+        }],
+    )
+
+
 class GetConsecutiveParkingLocationVisits(unittest.TestCase):
   """Tests for _get_consecutive_parking_location_visits."""