aboutsummaryrefslogtreecommitdiff
path: root/lib/GTFS.hs
blob: cf4c939360e09026bd47b8ce514625bf03b314fa (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE DeriveAnyClass             #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE DerivingStrategies         #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE NamedFieldPuns             #-}
{-# LANGUAGE RecordWildCards            #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE StandaloneKindSignatures   #-}
{-# LANGUAGE TupleSections              #-}
{-# LANGUAGE TypeApplications           #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE UndecidableInstances       #-}

-- | All kinds of stuff that has to deal with GTFS directly
-- (i.e. parsing, querying, Aeson instances, etc.)
module GTFS where

import qualified Codec.Archive.Zip              as Zip
import           Control.Monad.ST               (runST)
import           Data.Aeson                     (FromJSON,
                                                 Options (fieldLabelModifier),
                                                 ToJSON, defaultOptions,
                                                 genericParseJSON,
                                                 genericToJSON)
import qualified Data.Aeson                     as A
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)
import           Data.Text                      (Text)
import           Data.Time                      (UTCTime (utctDay), dayOfWeek,
                                                 getCurrentTime)
import           Data.Time.Calendar             (Day, DayOfWeek (..))
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                ((=~))
import           Web.Internal.FormUrlEncoded    (FromForm (..))
import           Web.Internal.HttpApiData       (FromHttpApiData (..))
-- import Data.Aeson.Generic (Options(fieldLabelModifier), deriveJSON, defaultOptions)
import           Control.Lens
import           Data.Char                      (toLower)
import           Data.Foldable                  (Foldable (fold))
import           Data.Map                       (Map)
import qualified Data.Map                       as M
import           Data.Proxy                     (Proxy (Proxy))
import           Data.Swagger                   (ParamSchema (..),
                                                 SchemaOptions,
                                                 ToSchema (declareNamedSchema),
                                                 defaultSchemaOptions,
                                                 genericDeclareNamedSchema)
import qualified Data.Swagger                   as S
import qualified Data.Text                      as T

aesonOptions prefix =
  defaultOptions { fieldLabelModifier = fieldModifier (T.length prefix) }
  where fieldModifier n label = case drop n label of
         c:rest -> toLower c : rest
         ""     -> ""

swaggerOptions :: Text -> SchemaOptions
swaggerOptions prefix =
  defaultSchemaOptions { S.fieldLabelModifier = fieldModifier (T.length prefix) }
  where fieldModifier n label = case drop n label of
         c:rest -> toLower c : rest
         ""     -> ""

newtype Time = Time { toSeconds :: Int }
  deriving newtype (ToJSON, FromJSON)
  deriving (Generic)

instance CSV.FromField Time where
  parseField f = do
    text :: String <- CSV.parseField f
    let (_,_,_,subs) = text =~ ("([0-9][0-9]?):([0-9][0-9]?):([0-9][0-9]?)" :: Text)
          :: (String, String, String, [String])
    case subs of
      [hh,mm,ss] -> pure $ Time $ read hh * 3600 + read mm * 60 + read ss
      _          -> fail $ "encountered an invalid date: " <> text

instance Show Time where
  show (Time seconds) = ""
    +|seconds `div` 3600|+":"
    +|(seconds `mod` 3600) `div` 60|+":"
    +|seconds `mod` 60|+""

instance CSV.FromField Day where
  parseField f = do
    text :: String <- CSV.parseField f
    let (_,_,_,subs) = text =~ ("([0-9][0-9][0-9][0-9])([0-9][0-9])([0-9][0-9])" :: Text)
          :: (String, String, String, [String])
    case subs of
      [yyyy,mm,dd] -> do
        let Just dayOfYear = monthAndDayToDayOfYearValid (Day.isLeapYear (read yyyy)) (read mm) (read dd)
        pure $ Day.fromOrdinalDate (read yyyy) dayOfYear
      _ -> fail $ "invalid date encountered: " <> show f

instance ToSchema Time where
  declareNamedSchema _ = do
    dings <- declareNamedSchema (Proxy @Int)
    pure (set (S.schema . S.description) (Just "Zeit in Sekunden seit Tagesanfang") dings)


data Depth = Shallow | Deep
type Switch :: Depth -> Type -> Type -> Type
type family Switch c a b where
  Switch Deep a b = a
  Switch Shallow a b = b
type family Optional c a where
  Optional Deep a = a
  Optional Shallow _ = ()

type StationID = Text
type TripID = Text
type ServiceID = Text


-- | This is what's called a Stop in GTFS
data Station = Station
  { stationId   :: StationID
  , stationName :: Text
  , stationLat  :: Float
  , stationLon  :: Float
  } deriving (Show, Generic)

instance ToSchema Station where
  declareNamedSchema = genericDeclareNamedSchema (swaggerOptions "station")
instance FromJSON Station where
  parseJSON = genericParseJSON (aesonOptions "station")
instance ToJSON Station where
  toJSON = genericToJSON (aesonOptions "station")


-- | This is what's called a stop time in GTFS
data Stop (deep :: Depth) = Stop
  { stopTrip      :: TripID
  , stopArrival   :: Time
  , stopDeparture :: Time
  , stopStation   :: Switch deep Station StationID
  , stopSequence  :: Int
  } deriving Generic

deriving instance Show (Stop 'Shallow)
deriving instance Show (Stop 'Deep)
instance FromJSON (Switch a Station StationID) => FromJSON (Stop a) where
  parseJSON = genericParseJSON (aesonOptions "stop")
instance ToJSON (Switch a Station StationID) => ToJSON (Stop a) where
  toJSON = genericToJSON (aesonOptions "stop")
instance ToSchema (Stop Deep)where
  declareNamedSchema = genericDeclareNamedSchema (swaggerOptions "stop")

data Calendar = Calendar
  { calServiceId :: Text
  , calMonday    :: Bool
  , calTuesday   :: Bool
  , calWednesday :: Bool
  , calThursday  :: Bool
  , calFriday    :: Bool
  , calSaturday  :: Bool
  , calSunday    :: Bool
  , calStartDate :: Day
  , calEndDate   :: Day
  } deriving (Show, Generic)



data CalendarExceptionType = ServiceAdded | ServiceRemoved
  deriving (Show, Eq, Generic, ToJSON, FromJSON)

instance FromHttpApiData CalendarExceptionType where
  parseUrlPiece = \case
    "added"   -> Right ServiceAdded
    "removed" -> Right ServiceRemoved
    unknown   -> Left ("unknown CalendarExceptionType: "<>unknown)

data CalendarDate = CalendarDate
  { caldateServiceId     :: Text
  , caldateDate          :: Day
  , caldateExceptionType :: CalendarExceptionType
  } deriving (Show, Generic)

instance FromJSON CalendarDate where
  parseJSON = genericParseJSON (aesonOptions "caldate")
instance ToJSON CalendarDate where
  toJSON = genericToJSON (aesonOptions "caldate")
instance FromForm CalendarDate

data Trip (deep :: Depth) (shape :: Depth)= Trip
  { tripRoute     :: Text
  , tripTripID    :: TripID
  , tripHeadsign  :: Maybe Text
  , tripShortName :: Maybe Text
  , tripDirection :: Maybe Bool
  -- NOTE: there's also block_id, which we're unlikely to ever need
  , tripServiceId :: Text
  -- , tripWheelchairAccessible :: Bool
  -- , tripBikesAllowed :: Bool
  , tripShape     :: Switch shape Shape Text
  , tripStops     :: Optional deep (Vector (Stop deep))
  } deriving Generic


deriving instance Show (Trip Shallow Shallow)
deriving instance Show (Trip Deep Deep)
deriving instance Show (Trip Deep Shallow)
instance (FromJSON (Switch d Shape Text), FromJSON (Optional d (Vector (Stop d))), FromJSON (Switch s Shape Text)) => FromJSON (Trip d s) where
  parseJSON = genericParseJSON (aesonOptions "trip")
instance (ToJSON (Switch d Shape Text), ToJSON (Optional d (Vector (Stop d))), ToJSON (Switch s Shape Text)) => ToJSON (Trip d s) where
  toJSON = genericToJSON (aesonOptions "trip")
instance ToSchema (Trip Deep Deep) where
  declareNamedSchema = genericDeclareNamedSchema (swaggerOptions "trip")
instance ToSchema (Trip Deep Shallow) where
  declareNamedSchema = genericDeclareNamedSchema (swaggerOptions "trip")

-- | helper function to find things in Vectors of things
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"
    <*> r .: "stop_name"
    <*> r .: "stop_lat"
    <*> r .: "stop_lon"

instance CSV.FromNamedRecord (Stop 'Shallow) where
  parseNamedRecord r = Stop
    <$> r .: "trip_id"
    <*> r .: "arrival_time"
    <*> r .: "departure_time"
    <*> r .: "stop_id"
    <*> r .: "stop_sequence"

instance CSV.FromNamedRecord Calendar where
  parseNamedRecord r = Calendar
    <$> r .: "service_id"
    <*> intAsBool' r "monday"
    <*> intAsBool' r "tuesday"
    <*> intAsBool' r "wednesday"
    <*> intAsBool' r "thursday"
    <*> intAsBool' r "friday"
    <*> intAsBool' r "saturday"
    <*> intAsBool' r "sunday"
    <*> 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
  pure $ case int :: Int of
    1 -> Just True
    0 -> Just False
    _ -> Nothing

intAsBool' :: CSV.NamedRecord -> BS.ByteString -> CSV.Parser Bool
intAsBool' r field = intAsBool r field >>= maybe
  (fail "unexpected value for a boolean.")
  pure


instance CSV.FromNamedRecord CalendarDate where
  parseNamedRecord r = CalendarDate
    <$> r .: "service_id"
    <*> r .: "date"
    <*> do
      int <- r .: "exception_type"
      case int :: Int of
        1 -> pure ServiceAdded
        2 -> pure ServiceRemoved
        _ -> fail $ "unexpected value in exception_type: "+|int|+"."


instance CSV.FromNamedRecord (Trip Shallow Shallow) where
  parseNamedRecord r = Trip
    <$> r .: "route_id"
    <*> r .: "trip_id"
    <*> r .: "trip_headsign"
    <*> r .: "trip_short_name"
    <*> intAsBool r "direction_id"
    <*> r .: "service_id"
    -- NOTE: these aren't booleans but triple-values
    -- <*> intAsBool r "wheelchair_accessible"
    -- <*> intAsBool r "bikes_allowed"
    <*> r .: "shape_id"
    <*> pure ()

data RawGTFS = RawGTFS
  { rawStations      :: Vector Station
  , rawStops         :: Vector (Stop Shallow)
  , rawTrips         :: Vector (Trip Shallow Shallow)
  , rawCalendar      :: Maybe (Vector Calendar)
  , rawCalendarDates :: Maybe (Vector CalendarDate)
  , rawShapePoints   :: Maybe (Vector ShapePoint)
  }


data GTFS = GTFS
  { stations      :: Map StationID Station
  , trips         :: Map TripID (Trip Deep Deep)
  , calendar      :: Map DayOfWeek (Vector Calendar)
  , calendarDates :: Map Day (Vector CalendarDate)
  , shapes        :: Map Text Shape

  , fancyCalendar :: Day -> (Vector ServiceID, Vector (Trip Deep Deep))
  -- ^ a more "fancy" encoding of the calendar?
  } -- deriving Show



loadRawGtfs :: FilePath -> IO RawGTFS
loadRawGtfs path = do
    zip <- Zip.toArchive <$> LB.readFile "./gtfs.zip"
    RawGTFS
      <$> decodeTable' "stops.txt" zip
      <*> decodeTable' "stop_times.txt" zip
      <*> 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 =
        case Zip.findEntryByPath path zip of
          Nothing -> pure Nothing
          Just csv -> case CSV.decodeByName (Zip.fromEntry csv) of
            Left err -> error $ "could not decode file "+|path|+": "+|err|+"."
            Right (_,v :: a) -> pure (Just v)
      decodeTable' path zip =
        decodeTable path zip >>= \case
          Nothing -> fail $ "required file "+|path|+" not found in gtfs.zip"
          Just a  -> pure a

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' shapes) rawTrips
    pure $ GTFS
      { stations =
        M.fromList $ (\station -> (stationId station, station))
        <$> V.toList rawStations
      , trips =
        M.fromList $ (\trip -> (tripTripID trip, trip))
        <$> V.toList trips'
      , calendar =
        fmap V.fromList
        $ M.fromListWith (<>)
        $ concatMap (\cal -> (, [cal]) <$> weekdays cal)
        $ V.toList (fromMaybe mempty rawCalendar)
      , calendarDates =
        fmap V.fromList
        $ M.fromListWith (<>) $ (\cd -> (caldateDate cd, [cd]))
        <$> V.toList (fromMaybe mempty rawCalendarDates)
      , shapes
      }
    where
      weekdays Calendar{..} =
        if calMonday then [Monday] else []
        <> if calTuesday then [Tuesday] else []
        <> if calWednesday then [Wednesday] else []
        <> if calThursday then [Thursday] else []
        <> if calFriday then [Friday] else []
        <> if calSaturday then [Saturday] else []
        <> [Sunday | calSunday]
      pushStop :: Vector Station -> Stop Shallow -> IO (Stop Deep)
      pushStop stations stop = do
        station <- case tableLookup stationId (stopStation stop) stations of
          Just a -> pure a
          Nothing -> fail $ "station with id "+|stopStation stop|+"is mentioned but not defined."
        pure $ stop { stopStation = station }
      pushTrip :: Vector (Stop Deep) -> Map Text Shape -> Trip Shallow Shallow -> IO (Trip Deep Deep)
      pushTrip stops shapes trip = if V.length alongRoute < 2
          then fail $ "trip with id "+|tripTripID trip|+" has no 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
servicesOnDay GTFS{..} day =
  fmap caldateServiceId added <> V.filter notCancelled regular
  where (added,removed) =
          V.partition (\cd -> caldateExceptionType cd == ServiceAdded)
          . fromMaybe mempty $ M.lookup day calendarDates
        regular = maybe mempty (fmap calServiceId) $ M.lookup (dayOfWeek day) calendar
        notCancelled serviceID =
          null (tableLookup caldateServiceId serviceID removed)

tripsOfService :: GTFS -> ServiceID -> Map TripID (Trip Deep Deep)
tripsOfService GTFS{..} serviceId =
  M.filter (\trip -> tripServiceId trip == serviceId ) trips

-- TODO: this should filter out trips ending there
tripsAtStation :: GTFS -> StationID -> Vector TripID
tripsAtStation GTFS{..} at = fmap stopTrip stops
  where
    stops = V.filter (\(stop :: Stop Deep) -> stationId (stopStation stop) == at) stops

tripsOnDay :: GTFS -> Day -> Map TripID (Trip Deep Deep)
tripsOnDay gtfs today = foldMap (tripsOfService gtfs) (servicesOnDay gtfs today)