Initial implementation of 'next arrivals'. Needs more work.

src/gtfs/mod.rs

@@ -1,11 +1,12 @@
 mod loader;
 mod utils;
 pub mod structs;
+use chrono::{DateTime, Local, NaiveDate, NaiveDateTime};
 use std::{
-    collections::{HashMap},
+    collections::{HashMap, HashSet},
     fs::File,
 };
-use gtfs_structures::{Agency, Calendar, CalendarDate, RawStopTime, RawTrip, Route, Stop, TimepointType};
+use gtfs_structures::{Agency, Calendar, CalendarDate, Exception, RawStopTime, RawTrip, Route, Stop, TimepointType, Trip};
 
 use crate::gtfs::{loader::load_gtfs, structs::{Arrival, Preferences}};
 
@@ -31,11 +32,64 @@ pub struct Gtfs {
 
 impl Gtfs {
 
-    pub fn _get_next_stops(&self) -> Vec<Arrival<'_>> {
-        let arrivals = Vec::<Arrival>::new();
+    pub fn get_next_arrivals_for(&self, target_datetime: &DateTime<Local>) -> Box<Vec<Arrival>> {
+        let naive_target = target_datetime.naive_local();
+        let target_date = naive_target.date();
+        
+        // Find which calendars apply
+        let mut active_service_ids: HashSet<String> = HashSet::new();
+        for (id, calendar) in self.calendar.iter() {
+            if calendar.valid_weekday(target_date) 
+            && calendar.start_date <= target_date 
+            && calendar.end_date > target_date {
+                active_service_ids.insert(id.to_string());
+            } 
+        }
 
-    return arrivals;
-}
+        // Are there any exceptions for the calendars above?
+        for (_calendar_id, exceptions) in self.calendar_dates.iter() {
+            for exception in exceptions.iter() {
+                if exception.date.eq(&target_date) {
+                    match exception.exception_type {
+                        Exception::Added =>  {
+                            active_service_ids.insert(exception.service_id.clone());
+                        }
+                        Exception::Deleted => {
+                            active_service_ids.remove(&exception.service_id);
+                        } 
+                    }
+                }
+            }                
+        }
+
+        // Find the trips happening on these calendars
+        let mut trips: HashMap<&String, &RawTrip> = HashMap::new();
+        for (id, trip) in self.trips.iter() {
+            if active_service_ids.contains(&trip.service_id) {
+                trips.insert(&trip.id, trip);
+            }
+        }
+
+        // Finally, find the arrivals for the active trips on the chosen stops
+        let mut arrivals: Vec<Arrival> = Vec::new();
+        let current_timestamp = target_datetime.timestamp();
+        for (id, stop_time) in self.stop_times.iter() {
+            if trips.contains_key(&stop_time.trip_id) {
+                let stop_timestamp = stop_time.departure_time.or(stop_time.arrival_time).unwrap();
+                if current_timestamp < stop_timestamp.into() {
+                    let arrival: Arrival = Arrival {
+                        route: self.routes.get(&self.trips.get(&stop_time.trip_id).unwrap().route_id).unwrap(),
+                        stop: self.stops.get(&stop_time.stop_id).unwrap(),
+                        departure_time: stop_timestamp
+                    };
+                    arrivals.push(arrival);                        
+                }
+            }
+        }
+
+        return Box::from(arrivals);
+    }
+        
 
     /// Load a GTFS structure from a zip file
     pub fn load(src_file: &str, prefs: &Preferences) -> Gtfs {
@@ -53,9 +107,8 @@ impl Gtfs {
         };
 
 
-    load_gtfs(&mut gtfs, &mut zip_reader, &prefs.route_numbers, &prefs.stop_codes);
+        load_gtfs(&mut gtfs, &mut zip_reader, &prefs.route_numbers, &prefs.stop_codes);
 
-    return gtfs;
+        return gtfs;
+    }
 }
-
-}