remove some extrapolation bugs

4 files changed, 66 insertions, 61 deletions

diff --git a/lib/Extrapolation.hs b/lib/Extrapolation.hs
index f505e73..5adc074 100644
--- a/lib/Extrapolation.hs
+++ b/lib/Extrapolation.hs
@@ -4,10 +4,11 @@
 {-# LANGUAGE DataKinds                  #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE LambdaCase                 #-}
+{-# LANGUAGE MultiWayIf                 #-}
 {-# LANGUAGE NamedFieldPuns             #-}
 {-# LANGUAGE RecordWildCards            #-}
 
-module Extrapolation (Extrapolator(..), LinearExtrapolator, linearDelay, distanceAlongLine) where
+module Extrapolation (Extrapolator(..), LinearExtrapolator(..), linearDelay, distanceAlongLine, euclid) where
 import           Data.Foldable      (maximumBy, minimumBy)
 import           Data.Function      (on)
 import           Data.List.NonEmpty (NonEmpty)
@@ -21,7 +22,8 @@ import           GHC.Float          (int2Double)
 import           GHC.IO             (unsafePerformIO)
 
 import           Conduit            (MonadIO (liftIO))
-import           Data.List          (sortBy)
+import           Data.List          (sortBy, sortOn)
+import           Data.Ord           (Down (..))
 import           GTFS               (Depth (Deep), GTFS (..), Seconds (..),
                                      Shape (..), Station (stationName),
                                      Stop (..), Time, Trip (..), seconds2Double,
@@ -30,38 +32,36 @@ import           Persist            (Running (..), TrainAnchor (..),
                                      TrainPing (..))
 import           Server.Util        (utcToSeconds)
 
+-- | Determines how to extrapolate delays (and potentially other things) from the real-time
+-- data sent in by the OBU. Potentially useful to swap out the algorithm, or give it options.
+-- TODO: maybe split into two classes?
 class Extrapolator a where
   -- | here's a position ping, guess things from that!
-  extrapolateAnchorFromPing :: GTFS -> Running -> TrainPing -> TrainAnchor
+  extrapolateAnchorFromPing :: a -> GTFS -> Running -> TrainPing -> TrainAnchor
 
   -- | extrapolate status at some time (i.e. "how much delay does the train have *now*?")
-  extrapolateAtSeconds :: NonEmpty TrainAnchor -> Seconds -> Maybe TrainAnchor
+  extrapolateAtSeconds :: a -> NonEmpty TrainAnchor -> Seconds -> Maybe TrainAnchor
   -- | extrapolate status at some places (i.e. "how much delay will it have at the next station?")
-  extrapolateAtPosition :: NonEmpty TrainAnchor -> Double -> Maybe TrainAnchor
+  extrapolateAtPosition :: a -> NonEmpty TrainAnchor -> Double -> Maybe TrainAnchor
 
-data LinearExtrapolator
+
+data LinearExtrapolator = LinearExtrapolator
 
 instance Extrapolator LinearExtrapolator where
-  -- TODO: this kind of sorting is bullshit; should look up depending on position,
-  -- not time (else there's glitches)
-  extrapolateAtSeconds history secondsNow =
+  extrapolateAtSeconds _ history secondsNow =
     fmap (minimumBy (compare `on` difference))
     $ NE.nonEmpty $ NE.filter (\a -> trainAnchorWhen a < secondsNow) history
     where difference status = secondsNow - (trainAnchorWhen status)
 
   -- note that this sorts (descending) for time first as a tie-breaker
   -- (in case the train just stands still for a while, take the most recent update)
-  extrapolateAtPosition history positionNow =
+  extrapolateAtPosition _ history positionNow =
     fmap (minimumBy (compare `on` difference))
-    $ NE.nonEmpty $ sortBy (flippedCompare `on` trainAnchorWhen)
+    $ NE.nonEmpty $ sortOn (Down . trainAnchorWhen)
     $ NE.filter (\a -> trainAnchorSequence a < positionNow) history
     where difference status = positionNow - (trainAnchorSequence status)
-          flippedCompare a b = case compare a b of
-            LT -> GT
-            GT -> LT
-            a  -> a
 
-  extrapolateAnchorFromPing gtfs@GTFS{..} Running{..} ping@TrainPing{..} = TrainAnchor
+  extrapolateAnchorFromPing _ gtfs@GTFS{..} Running{..} ping@TrainPing{..} = TrainAnchor
     { trainAnchorCreated = trainPingTimestamp
     , trainAnchorTrip = runningTrip
     , trainAnchorDay = runningDay
@@ -74,44 +74,49 @@ instance Extrapolator LinearExtrapolator where
           (trainAnchorDelay, trainAnchorSequence) = linearDelay gtfs trip ping runningDay
 
 linearDelay :: GTFS -> Trip Deep Deep -> TrainPing -> Day -> (Seconds, Double)
-linearDelay GTFS{..} trip@Trip{..} TrainPing{..} runningDay = unsafePerformIO $ do
-  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
-        expectedProgress = crop $
-            seconds2Double (utcToSeconds trainPingTimestamp runningDay
-                         - toSeconds (stopDeparture lastStop) tzseries runningDay)
-            / seconds2Double expectedTravelTime
-            where crop a
-                    | a < 0 = 0
-                    | a > 1 = 1
-                    | otherwise = a
-        observedProgress =
-            distanceAlongLine line (stationGeopos $ stopStation lastStop) closestPoint
-            / distanceAlongLine line (stationGeopos $ stopStation lastStop) (stationGeopos $ stopStation nextStop)
-        observedPosition =
-          (int2Double $ stopSequence lastStop) + observedProgress * (int2Double $ stopSequence nextStop - stopSequence lastStop)
+linearDelay GTFS{..} trip@Trip{..} TrainPing{..} runningDay = (observedDelay, observedSequence)
+  where -- | at which sequence number is the ping?
+        observedSequence = int2Double (stopSequence lastStop)
+          + observedProgress * int2Double (stopSequence nextStop - stopSequence lastStop)
+        -- | how much later/earlier is the ping than would be expected?
         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
+          + if expectedProgress /= 1 then 0
+            else seconds2Double (utcToSeconds trainPingTimestamp runningDay
+                                 - toSeconds (stopArrival nextStop) tzseries runningDay)
+
+        -- | how far along towards the next station is the ping (between 0 and 1)?
+        observedProgress =
+            distanceAlongLine line (stationGeopos $ stopStation lastStop) closestPoint
+            / distanceAlongLine line (stationGeopos $ stopStation lastStop) (stationGeopos $ stopStation nextStop)
+        -- | to compare: where would a linearly-moving train be (between 0 and 1)?
+        expectedProgress = if
+          | p < 0     -> 0
+          | p > 1     -> 1
+          | otherwise -> p
+          where p = seconds2Double (utcToSeconds trainPingTimestamp runningDay
+                                    - toSeconds (stopDeparture lastStop) tzseries runningDay)
+                    / seconds2Double expectedTravelTime
+        -- | how long do we expect the trip from last to next station to take?
+        expectedTravelTime =
+            toSeconds (stopArrival nextStop) tzseries runningDay
+            - toSeconds (stopDeparture lastStop) tzseries runningDay
+
+        closestPoint = minimumBy (compare `on` euclid (trainPingLat, trainPingLong)) line
+        line = shapePoints tripShape
+        lastStop = tripStops V.! (nextIndex - 1)
+        nextStop = tripStops V.! nextIndex
+        -- | index of the /next/ stop in the list, except when we're already at the last stop
+        -- (in which case it stays the same)
+        nextIndex = if
+          | null remaining -> length tripStops - 1
+          | idx' == 0      -> 1
+          | otherwise      -> idx'
+          where idx' = fst $ V.minimumBy (compare `on` snd) remaining
+                remaining = V.filter (\(_,dist) -> dist > 0) $ V.indexed
+                  $ fmap (distanceAlongLine line closestPoint . stationGeopos . stopStation) tripStops
 
 distanceAlongLine :: V.Vector (Double, Double) -> (Double, Double) -> (Double, Double) -> Double
 distanceAlongLine line p1 p2 = along2 - along1
@@ -125,8 +130,8 @@ distanceAlongLine line p1 p2 = along2 - along1
         sumSegments line = snd
           $ foldl (\(p,a) p' -> (p', a + euclid p p')) (V.head line,0) $ line
 
-
-
+-- | euclidean distance. Notably not applicable when you're on a sphere
+-- (but good enough when the sphere is the earth)
 euclid :: Floating f => (f,f) -> (f,f) -> f
 euclid (x1,y1) (x2,y2) = sqrt (x*x + y*y)
   where x = x1 - x2
diff --git a/lib/Server.hs b/lib/Server.hs
index 3c1e84b..93046f8 100644
--- a/lib/Server.hs
+++ b/lib/Server.hs
@@ -53,7 +53,7 @@ import           Yesod                       (toWaiAppPlain)
 
 import           Conferer                    (fetch, mkConfig)
 import           Extrapolation               (Extrapolator (..),
-                                              LinearExtrapolator)
+                                              LinearExtrapolator (..))
 import           System.IO.Unsafe
 
 import           Config                      (ServerConfig)
@@ -105,7 +105,7 @@ server gtfs@GTFS{..} dbpool = handleDebugAPI
             onError
             pure Nothing
           Just running@Running{..} -> do
-            let anchor = extrapolateAnchorFromPing @LinearExtrapolator gtfs running ping
+            let anchor = extrapolateAnchorFromPing LinearExtrapolator gtfs running ping
             -- TODO: are these always inserted in order?
             runSql dbpool $ do
               insert ping
@@ -129,7 +129,7 @@ server gtfs@GTFS{..} dbpool = handleDebugAPI
                 -- ignore this and continue sending messages, which will continue to be handled.
                 liftIO $ handleTrainPing (WS.sendClose conn ("" :: ByteString)) ping >>= \case
                   Just anchor -> WS.sendTextData conn (A.encode anchor)
-                  Nothing -> pure ()
+                  Nothing     -> pure ()
         handleDebugState = do
           now <- liftIO getCurrentTime
           runSql dbpool $ do
diff --git a/lib/Server/ControlRoom.hs b/lib/Server/ControlRoom.hs
index e5e9b7c..b1948f2 100644
--- a/lib/Server/ControlRoom.hs
+++ b/lib/Server/ControlRoom.hs
@@ -53,7 +53,7 @@ import           Yesod
 import           Yesod.Form
 
 import           Extrapolation            (Extrapolator (..),
-                                           LinearExtrapolator)
+                                           LinearExtrapolator (..))
 import           GTFS
 import           Numeric                  (showFFloat)
 import           Persist
@@ -234,8 +234,8 @@ getTrainViewR trip day = do
         $else
           <td title="_{MsgBlockToken}"><a href="@{TokenBlock key}">_{MsgBlockToken}</a>
 |]
-          where guessDelay history = fmap trainAnchorDelay . extrapolateAtPosition @LinearExtrapolator history
-                guessAtSeconds = extrapolateAtSeconds @LinearExtrapolator
+          where guessDelay history = fmap trainAnchorDelay . extrapolateAtPosition LinearExtrapolator history
+                guessAtSeconds = extrapolateAtSeconds LinearExtrapolator
 
 
 getTripsViewR :: Handler Html
diff --git a/lib/Server/GTFS_RT.hs b/lib/Server/GTFS_RT.hs
index e0bbeda..1641131 100644
--- a/lib/Server/GTFS_RT.hs
+++ b/lib/Server/GTFS_RT.hs
@@ -77,7 +77,7 @@ import           Data.UUID                                  (toASCIIBytes,
                                                              toLazyASCIIBytes)
 import qualified Data.Vector                                as V
 import           Extrapolation                              (Extrapolator (extrapolateAtPosition, extrapolateAtSeconds),
-                                                             LinearExtrapolator)
+                                                             LinearExtrapolator (..))
 import           GTFS                                       (Depth (..))
 import           GTFS.Realtime.TripUpdate                   (TripUpdate (TripUpdate))
 import           Server.Util                                (Service,
@@ -149,9 +149,9 @@ gtfsRealtimeServer gtfs@GTFS{..} dbpool = handleServiceAlerts :<|> handleTripUpd
 
           dFeedMessage $ Seq.fromList $ mkTripUpdate nowSeconds <$> anchors
           where mkTripUpdate nowSeconds (Entity (RunningKey (Token uuid)) Running{..}, Trip{..} :: Trip Deep Deep, anchors) =
-                  let lastCall = extrapolateAtSeconds @LinearExtrapolator anchors nowSeconds
+                  let lastCall = extrapolateAtSeconds LinearExtrapolator anchors nowSeconds
                       stations = tripStops
-                        <&> (\stop@Stop{..} -> fmap (, stop) $ extrapolateAtPosition @LinearExtrapolator anchors (int2Double stopSequence))
+                        <&> (\stop@Stop{..} -> fmap (, stop) $ extrapolateAtPosition LinearExtrapolator anchors (int2Double stopSequence))
                   in (dFeedEntity (Utf8 $ toLazyASCIIBytes uuid))
                   { FE.trip_update = Just $ TripUpdate
                     { TU.trip = dTripDescriptor runningTrip (Just runningDay)