From 3e27228d00efef1bbcfd4700a4de32546a244156 Mon Sep 17 00:00:00 2001
From: Walton Hoops <me@waltonhoops.com>
Date: Fri, 15 Nov 2024 14:59:12 -0700
Subject: [PATCH] feat: use last predicted stop as added trip headsign

---
 apps/state/lib/state/trip/added.ex        | 48 ++++++-----------------
 apps/state/test/state/trip/added_test.exs | 16 ++++----
 2 files changed, 20 insertions(+), 44 deletions(-)

diff --git a/apps/state/lib/state/trip/added.ex b/apps/state/lib/state/trip/added.ex
index ee9ffb20..eab50b53 100644
--- a/apps/state/lib/state/trip/added.ex
+++ b/apps/state/lib/state/trip/added.ex
@@ -56,6 +56,10 @@ defmodule State.Trip.Added do
     with %{route_pattern_id: route_pattern_id} when is_binary(route_pattern_id) <- prediction,
          %{representative_trip_id: rep_trip_id} <- State.RoutePattern.by_id(route_pattern_id),
          [trip | _] <- State.Trip.by_id(rep_trip_id) do
+      stop = parent_or_stop(prediction.stop_id)
+
+      headsign = if stop, do: stop.name, else: trip.headsign
+
       [
         %{
           trip
@@ -64,7 +68,8 @@ defmodule State.Trip.Added do
             service_id: nil,
             wheelchair_accessible: 1,
             bikes_allowed: 0,
-            revenue: prediction.revenue
+            revenue: prediction.revenue,
+            headsign: headsign
         }
       ]
     else
@@ -74,25 +79,7 @@ defmodule State.Trip.Added do
   end
 
   defp prediction_to_trip_via_shape(prediction) do
-    stop =
-      case State.Stop.by_id(prediction.stop_id) do
-        %{parent_station: nil} = stop -> stop
-        %{parent_station: id} -> State.Stop.by_id(id)
-        _other -> nil
-      end
-
-    last_stop_id =
-      [prediction.route_id]
-      |> State.Shape.select_routes(prediction.direction_id)
-      |> Stream.filter(&(&1.route_id == prediction.route_id))
-      |> Enum.find_value(&last_stop_id_on_shape(&1, prediction, stop))
-
-    stop =
-      if is_nil(last_stop_id) or last_stop_id == stop.id do
-        stop
-      else
-        State.Stop.by_id(last_stop_id)
-      end
+    stop = parent_or_stop(prediction.stop_id)
 
     if stop == nil do
       []
@@ -116,22 +103,11 @@ defmodule State.Trip.Added do
     end
   end
 
-  defp last_stop_id_on_shape(_, _, nil), do: nil
-
-  defp last_stop_id_on_shape(%{priority: p} = shape, prediction, stop) when p >= 0 do
-    shape_stops =
-      State.StopsOnRoute.by_route_id(
-        prediction.route_id,
-        direction_id: prediction.direction_id,
-        shape_ids: [shape.id]
-      )
-
-    if Enum.any?(shape_stops, &(&1 in [stop.id, stop.parent_station])) do
-      List.last(shape_stops)
+  defp parent_or_stop(stop_id) do
+    case State.Stop.by_id(stop_id) do
+      %{parent_station: nil} = stop -> stop
+      %{parent_station: id} -> State.Stop.by_id(id)
+      _other -> nil
     end
   end
-
-  defp last_stop_id_on_shape(_, _, _) do
-    nil
-  end
 end
diff --git a/apps/state/test/state/trip/added_test.exs b/apps/state/test/state/trip/added_test.exs
index cd21f273..fd8abd91 100644
--- a/apps/state/test/state/trip/added_test.exs
+++ b/apps/state/test/state/trip/added_test.exs
@@ -181,10 +181,10 @@ defmodule State.Trip.AddedTest do
       {:ok, %{shape: shape, prediction: prediction}}
     end
 
-    test "if there's a matching shape for the route/direction, uses the last stop from that shape",
+    test "if there's a matching shape for the route/direction, uses the last stop predicted stop",
          %{prediction: prediction} do
       insert_predictions([prediction])
-      assert [%{headsign: "Last Stop on Shape"}] = by_id(@trip_id)
+      assert [%{headsign: "Parent"}] = by_id(@trip_id)
     end
 
     test "does not use a shape on the wrong route or direction", %{
@@ -231,13 +231,13 @@ defmodule State.Trip.AddedTest do
     test "creates a trip with revenue value set to :REVENUE",
          %{prediction: prediction} do
       insert_predictions([%{prediction | revenue: :REVENUE}])
-      assert [%{headsign: "Last Stop on Shape", revenue: :REVENUE}] = by_id(@trip_id)
+      assert [%{headsign: "Parent", revenue: :REVENUE}] = by_id(@trip_id)
     end
 
     test "creates a trip with revenue value set to false",
          %{prediction: prediction} do
       insert_predictions([%{prediction | revenue: :NON_REVENUE}])
-      assert [%{headsign: "Last Stop on Shape", revenue: :NON_REVENUE}] = by_id(@trip_id)
+      assert [%{headsign: "Parent", revenue: :NON_REVENUE}] = by_id(@trip_id)
     end
   end
 
@@ -267,19 +267,19 @@ defmodule State.Trip.AddedTest do
       end)
     end
 
-    test "if there's a matching route_pattern, use the representative trip" do
+    test "if there's a matching route_pattern, use the last predicted stop" do
       insert_predictions([%{@prediction | stop_id: "child"}])
-      assert [%{headsign: "Headsign"}] = by_id(@trip_id)
+      assert [%{headsign: "Parent"}] = by_id(@trip_id)
     end
 
     test "creates a trip with revenue value set to :REVENUE" do
       insert_predictions([%{@prediction | stop_id: "child", revenue: :REVENUE}])
-      assert [%{headsign: "Headsign", revenue: :REVENUE}] = by_id(@trip_id)
+      assert [%{headsign: "Parent", revenue: :REVENUE}] = by_id(@trip_id)
     end
 
     test "creates a trip with revenue value set to :NON_REVENUE" do
       insert_predictions([%{@prediction | stop_id: "child", revenue: :NON_REVENUE}])
-      assert [%{headsign: "Headsign", revenue: :NON_REVENUE}] = by_id(@trip_id)
+      assert [%{headsign: "Parent", revenue: :NON_REVENUE}] = by_id(@trip_id)
     end
   end
 end