handle GTFS shapes (& send them out along trips)

unfortunately, this makes some responses pretty large. Let's see if it
works this way or if i have to find some way to break them up instead.

2 files changed, 64 insertions, 14 deletions

diff --git a/lib/GTFS.hs b/lib/GTFS.hs
index be80745..bd29b6d 100644
--- a/lib/GTFS.hs
+++ b/lib/GTFS.hs
@@ -6,6 +6,7 @@
 {-# LANGUAGE FlexibleInstances          #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE LambdaCase                 #-}
+{-# LANGUAGE NamedFieldPuns             #-}
 {-# LANGUAGE RecordWildCards            #-}
 {-# LANGUAGE StandaloneDeriving         #-}
 {-# LANGUAGE StandaloneKindSignatures   #-}
@@ -20,6 +21,7 @@ module GTFS where
 
 
 import qualified Codec.Archive.Zip              as Zip
+import           Control.Monad.ST               (runST)
 import           Data.Aeson                     (FromJSON,
                                                  Options (fieldLabelModifier),
                                                  ToJSON, defaultOptions,
@@ -30,6 +32,7 @@ import qualified Data.ByteString                as BS
 import qualified Data.ByteString.Lazy           as LB
 import           Data.Csv                       ((.:))
 import qualified Data.Csv                       as CSV
+import           Data.Function                  (on)
 import           Data.Functor                   ((<&>))
 import           Data.Kind                      (Type)
 import           Data.Maybe                     (fromJust, fromMaybe)
@@ -41,6 +44,7 @@ import           Data.Time.Calendar.MonthDay    (monthAndDayToDayOfYearValid)
 import qualified Data.Time.Calendar.OrdinalDate as Day
 import           Data.Vector                    (Vector)
 import qualified Data.Vector                    as V
+import qualified Data.Vector.Algorithms.Intro   as V
 import           Fmt                            ((+|), (|+))
 import           GHC.Generics                   (Generic)
 import           Text.Regex.TDFA                ((=~))
@@ -195,16 +199,16 @@ data Trip (deep :: Depth) = Trip
   , tripServiceId :: Text
   -- , tripWheelchairAccessible :: Bool
   -- , tripBikesAllowed :: Bool
-  , tripShapeId   :: Text
+  , tripShape     :: Switch deep Shape Text
   , tripStops     :: Optional deep (Vector (Stop deep))
   } deriving Generic
 
 
 deriving instance Show (Trip Shallow)
 deriving instance Show (Trip Deep)
-instance FromJSON (Optional d (Vector (Stop d))) => FromJSON (Trip d) where
+instance (FromJSON (Switch d Shape Text), FromJSON (Optional d (Vector (Stop d)))) => FromJSON (Trip d) where
   parseJSON = genericParseJSON (aesonOptions "trip")
-instance ToJSON (Optional d (Vector (Stop d))) => ToJSON (Trip d) where
+instance (ToJSON (Switch d Shape Text), ToJSON (Optional d (Vector (Stop d)))) => ToJSON (Trip d) where
   toJSON = genericToJSON (aesonOptions "trip")
 instance ToSchema (Trip Deep) where
   declareNamedSchema = genericDeclareNamedSchema (swaggerOptions "trip")
@@ -213,6 +217,25 @@ instance ToSchema (Trip Deep) where
 tableLookup :: Eq key => (a -> key) -> key -> Vector a -> Maybe a
 tableLookup proj key = V.find (\a -> proj a == key)
 
+data ShapePoint = ShapePoint
+  { shapePtId       :: Text
+  , shapePtLat      :: Double
+  , shapePtLong     :: Double
+  , shapePtSequence :: Int
+  } deriving Generic
+
+data Shape = Shape
+  { shapeId     :: Text
+  , shapePoints :: Vector (Double,Double)
+  } deriving (Generic, Show)
+
+instance FromJSON Shape where
+  parseJSON = genericParseJSON (aesonOptions "shape")
+instance ToJSON Shape where
+  toJSON = genericToJSON (aesonOptions "shape")
+instance ToSchema Shape where
+  declareNamedSchema = genericDeclareNamedSchema (swaggerOptions "shape")
+
 instance CSV.FromNamedRecord Station where
   parseNamedRecord r = Station
     <$> r .: "stop_id"
@@ -241,6 +264,13 @@ instance CSV.FromNamedRecord Calendar where
     <*> r .: "start_date"
     <*> r .: "end_date"
 
+instance CSV.FromNamedRecord ShapePoint where
+  parseNamedRecord r = ShapePoint
+    <$> r .: "shape_id"
+    <*> r .: "shape_pt_lat"
+    <*> r .: "shape_pt_lon"
+    <*> r .: "shape_pt_sequence"
+
 intAsBool :: CSV.NamedRecord -> BS.ByteString -> CSV.Parser (Maybe Bool)
 intAsBool r field = do
   int <- r .: field
@@ -287,6 +317,7 @@ data RawGTFS = RawGTFS
   , rawTrips         :: Vector (Trip Shallow)
   , rawCalendar      :: Maybe (Vector Calendar)
   , rawCalendarDates :: Maybe (Vector CalendarDate)
+  , rawShapePoints   :: Maybe (Vector ShapePoint)
   }
 
 
@@ -295,6 +326,7 @@ data GTFS = GTFS
   , trips         :: Map TripID (Trip Deep)
   , calendar      :: Map DayOfWeek (Vector Calendar)
   , calendarDates :: Map Day (Vector CalendarDate)
+  , shapes        :: Map Text Shape
 
   , fancyCalendar :: Day -> (Vector ServiceID, Vector (Trip Deep))
   -- ^ a more "fancy" encoding of the calendar?
@@ -311,6 +343,7 @@ loadRawGtfs path = do
       <*> decodeTable' "trips.txt" zip
       <*> decodeTable "calendar.txt" zip
       <*> decodeTable "calendar_dates.txt" zip
+      <*> decodeTable "shapes.txt" zip
     where
       decodeTable :: CSV.FromNamedRecord a => FilePath -> Zip.Archive -> IO (Maybe (Vector a))
       decodeTable path zip =
@@ -327,8 +360,13 @@ loadRawGtfs path = do
 loadGtfs :: FilePath -> IO GTFS
 loadGtfs path = do
     shallow@RawGTFS{..} <- loadRawGtfs path
+    -- TODO: sort these according to sequence numbers
+    let shapes =
+          V.foldr' sortShapePoint mempty
+          $ V.modify (V.sortBy (compare `on` shapePtSequence))
+          (fromMaybe mempty rawShapePoints)
     stops' <- V.mapM (pushStop rawStations) rawStops
-    trips' <- V.mapM (pushTrip stops') rawTrips
+    trips' <- V.mapM (pushTrip stops' shapes) rawTrips
     pure $ GTFS
       { stations =
         M.fromList $ (\station -> (stationId station, station))
@@ -345,6 +383,7 @@ loadGtfs path = do
         fmap V.fromList
         $ M.fromListWith (<>) $ (\cd -> (caldateDate cd, [cd]))
         <$> V.toList (fromMaybe mempty rawCalendarDates)
+      , shapes
       }
     where
       weekdays Calendar{..} =
@@ -361,13 +400,24 @@ loadGtfs path = do
           Just a -> pure a
           Nothing -> fail $ "station with id "+|stopStation stop|+"is mentioned but not defined."
         pure $ stop { stopStation = station }
-      pushTrip :: Vector (Stop Deep) -> Trip Shallow -> IO (Trip Deep)
-      pushTrip stops trip = if V.length alongRoute < 2
+      pushTrip :: Vector (Stop Deep) -> Map Text Shape -> Trip Shallow -> IO (Trip Deep)
+      pushTrip stops shapes trip = if V.length alongRoute < 2
           then fail $ "trip with id "+|tripTripID trip|+" has no stops"
-          else pure $ trip { tripStops = alongRoute }
-        where alongRoute = -- TODO: sort these according to stops
-                V.filter (\s -> stopTrip s == tripTripID trip) stops
-
+          else do
+            a <- case M.lookup (tripShape trip) shapes of
+              Nothing -> fail $ "trip with id "+|tripTripID trip|+" mentions a shape that does not exist."
+              Just a -> pure a
+            pure $ trip { tripStops = alongRoute, tripShape = a }
+        where alongRoute =
+                V.modify (V.sortBy (compare `on` stopSequence))
+                $ V.filter (\s -> stopTrip s == tripTripID trip) stops
+      sortShapePoint :: ShapePoint -> Map Text Shape -> Map Text Shape
+      sortShapePoint ShapePoint{..} shapes = M.alter appendPoint shapePtId shapes
+        where
+          point = (shapePtLat, shapePtLong)
+          appendPoint = \case
+            Just shape -> Just $ shape { shapePoints = V.cons point (shapePoints shape) }
+            Nothing -> Just $ Shape { shapeId = shapePtId, shapePoints = V.singleton point }
 
 
 servicesOnDay :: GTFS -> Day -> Vector ServiceID
diff --git a/lib/Persist.hs b/lib/Persist.hs
index 9486a1d..4a6d9b4 100644
--- a/lib/Persist.hs
+++ b/lib/Persist.hs
@@ -37,15 +37,15 @@ import           Control.Monad.Logger        (NoLoggingT)
 import           Control.Monad.Reader        (ReaderT)
 import           Data.Data                   (Proxy (..))
 import           Data.Pool                   (Pool)
-import           Data.Time                   (NominalDiffTime,
+import           Data.Time                   (NominalDiffTime, TimeOfDay,
                                               UTCTime (utctDay), addUTCTime,
                                               dayOfWeek, diffUTCTime,
-                                              getCurrentTime, nominalDay, TimeOfDay)
+                                              getCurrentTime, nominalDay)
+import           Data.Time.Calendar          (Day, DayOfWeek (..))
+import           Data.Vector                 (Vector)
 import           Database.Persist.Postgresql (SqlBackend)
 import           GHC.Generics                (Generic)
 import           Web.PathPieces              (PathPiece)
-import Data.Vector (Vector)
-import           Data.Time.Calendar             (Day, DayOfWeek (..))
 
 newtype Token = Token UUID
   deriving newtype