fix the close-to-a-station bug

(previously tracktrain could end up in a situation where the next and
last station weren't actually adjacent stops, which messed up the
prediction)

1 files changed, 31 insertions, 22 deletions

diff --git a/lib/Extrapolation.hs b/lib/Extrapolation.hs
index 770d4ce..cc77a92 100644
--- a/lib/Extrapolation.hs
+++ b/lib/Extrapolation.hs
@@ -7,7 +7,7 @@
 {-# LANGUAGE NamedFieldPuns             #-}
 {-# LANGUAGE RecordWildCards            #-}
 
-module Extrapolation (Extrapolator(..), LinearExtrapolator, linearDelay, secondsNow) where
+module Extrapolation (Extrapolator(..), LinearExtrapolator, linearDelay, secondsNow, distanceAlongLine) where
 import           Data.Foldable      (maximumBy, minimumBy)
 import           Data.Function      (on)
 import           Data.List.NonEmpty (NonEmpty)
@@ -24,7 +24,7 @@ import           Conduit            (MonadIO (liftIO))
 import           Data.List          (sortBy)
 import           GTFS               (Depth (Deep), GTFS (..), Seconds (..),
                                      Shape (..), Stop (..), Time, Trip (..),
-                                     seconds2Double, stationGeopos, toSeconds)
+                                     seconds2Double, stationGeopos, toSeconds, Station (stationName))
 import           Persist            (Running (..), TrainAnchor (..),
                                      TrainPing (..))
 
@@ -73,20 +73,20 @@ instance Extrapolator LinearExtrapolator where
 
 linearDelay :: GTFS -> Trip Deep Deep -> TrainPing -> Day -> (Seconds, Double)
 linearDelay GTFS{..} trip@Trip{..} TrainPing{..} runningDay = unsafePerformIO $ do
-  print (scheduledPosition, round $ (expectedProgress - observedProgress) * int2Double (unSeconds expectedTravelTime))
-  pure $ (Seconds $ round $ (expectedProgress - observedProgress) * int2Double (unSeconds expectedTravelTime)
-         , scheduledPosition)
-  where closestPoint =
-            minimumBy (compare `on` euclid (trainPingLat, trainPingLong)) line
-        nextStop = snd $
-            minimumBy (compare `on` fst)
-            $ V.filter (\(dist,_) -> dist > 0)
-            $ fmap (\stop -> (distanceAlongLine line closestPoint (stationGeopos $ stopStation stop), stop)) tripStops
-        lastStop = snd $
-            maximumBy (compare `on` fst)
-            $ V.filter (\(dist,_) -> dist < 0)
-            $ fmap (\stop -> (distanceAlongLine line closestPoint (stationGeopos $ stopStation stop), stop)) tripStops
+  print (observedPosition, observedProgress)
+  print (stationName . stopStation $ lastStop, stationName . stopStation $ nextStop)
+  print (distanceAlongLine line (stationGeopos $ stopStation lastStop) closestPoint
+        , distanceAlongLine line (stationGeopos $ stopStation lastStop) (stationGeopos $ stopStation nextStop))
+  pure (observedDelay, observedPosition)
+  where closestPoint = minimumBy (compare `on` euclid (trainPingLat, trainPingLong)) line
         line = shapePoints tripShape
+        lastStop = tripStops V.! (nextIndex - 1)
+        nextStop = tripStops V.! nextIndex
+        nextIndex = if idx' == 0
+          then 1 else idx'
+          where idx' = fst $ V.minimumBy (compare `on` snd)
+                  $ V.filter (\(_,dist) -> dist > 0) $ V.indexed
+                  $ fmap (distanceAlongLine line closestPoint . stationGeopos . stopStation) tripStops
         expectedTravelTime =
             toSeconds (stopArrival nextStop) tzseries runningDay
             - toSeconds (stopDeparture lastStop) tzseries runningDay
@@ -101,18 +101,27 @@ linearDelay GTFS{..} trip@Trip{..} TrainPing{..} runningDay = unsafePerformIO $
         observedProgress =
             distanceAlongLine line (stationGeopos $ stopStation lastStop) closestPoint
             / distanceAlongLine line (stationGeopos $ stopStation lastStop) (stationGeopos $ stopStation nextStop)
-        scheduledPosition =
+        observedPosition =
           (int2Double $ stopSequence lastStop) + observedProgress * (int2Double $ stopSequence nextStop - stopSequence lastStop)
+        observedDelay = Seconds $ round $
+          (expectedProgress - observedProgress) * int2Double (unSeconds expectedTravelTime)
+          -- if the hypothetical on-time train is already at (or past) the next station,
+          -- just add the time distance we're behind
+          + if expectedProgress == 1
+            then seconds2Double (utcToSeconds trainPingTimestamp runningDay - toSeconds (stopArrival nextStop) tzseries runningDay)
+            else 0
 
 distanceAlongLine :: V.Vector (Double, Double) -> (Double, Double) -> (Double, Double) -> Double
 distanceAlongLine line p1 p2 = along2 - along1
   where along1 = along p1
         along2 = along p2
-        along p@(x,y) = snd
-          $ foldl (\(p,a) p' -> (p', a + euclid p p')) (V.head line,0)
+        along p@(x,y) =
+          sumSegments
           $ V.take (index + 1) line
-          where index = fst $ minimumBy (compare `on` (euclid p . snd))
-                  $ V.indexed line
+          where index = V.minIndexBy (compare `on` euclid p) line
+        sumSegments :: V.Vector (Double, Double) -> Double
+        sumSegments line = snd
+          $ foldl (\(p,a) p' -> (p', a + euclid p p')) (V.head line,0) $ line
 
 -- | convert utc time to seconds on a day, with wrap-around
 -- for trains that cross midnight.
@@ -125,7 +134,7 @@ secondsNow runningDay = do
   now <- liftIO getCurrentTime
   pure $ utcToSeconds now runningDay
 
-euclid :: Fractional f => (f,f) -> (f,f) -> f
-euclid (x1,y1) (x2,y2) = x*x + y*y
+euclid :: Floating f => (f,f) -> (f,f) -> f
+euclid (x1,y1) (x2,y2) = sqrt (x*x + y*y)
   where x = x1 - x2
         y = y1 - y2