open-standards/map/map-rs/crates/map-route/src/transit.rs

//! Public-transit routing — earliest arrival over a timetable via the
//! **Connection Scan Algorithm** (CSA).
//!
//! Where [`crate::find_route`] and [`crate::cost`] route a continuous road graph,
//! transit is *time-dependent*: you can only board a vehicle at its scheduled
//! departure. A [`Timetable`] is a list of [`Connection`]s (one scheduled hop
//! between consecutive stops) plus walking [`footpaths`] for transfers. CSA scans
//! the connections once in departure order, relaxing the earliest arrival at each
//! stop — `O(connections)` per query after a sort, and exact for earliest arrival.
//!
//! Times are integer seconds (e.g. since service midnight), matching GTFS.
//!
//! [`footpaths`]: Timetable::add_footpath

use std::collections::BTreeMap;

use serde::{Deserialize, Serialize};

use crate::{haversine_distance, Coordinate};

const INF: i64 = i64::MAX;

/// One scheduled vehicle hop between two consecutive stops on a trip.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct Connection {
    pub trip: u64,
    pub dep_stop: u64,
    pub arr_stop: u64,
    pub dep_time: i64,
    pub arr_time: i64,
}

/// A transit network: stops, scheduled connections, and walking transfers.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Timetable {
    pub stops: BTreeMap<u64, Coordinate>,
    pub connections: Vec<Connection>,
    /// `stop -> [(reachable stop, walk seconds)]`.
    footpaths: BTreeMap<u64, Vec<(u64, i64)>>,
}

impl Timetable {
    pub fn add_stop(&mut self, id: u64, coord: Coordinate) {
        self.stops.insert(id, coord);
    }

    pub fn add_connection(&mut self, c: Connection) {
        self.connections.push(c);
    }

    /// Add a bidirectional walking transfer between two stops.
    pub fn add_footpath(&mut self, a: u64, b: u64, secs: i64) {
        self.footpaths.entry(a).or_default().push((b, secs));
        self.footpaths.entry(b).or_default().push((a, secs));
    }

    /// Generate transfers between every pair of stops within `max_dist_m`,
    /// timed at `walk_speed_mps` (e.g. 1.4). Replaces existing footpaths.
    pub fn auto_footpaths(&mut self, max_dist_m: f64, walk_speed_mps: f64) {
        self.footpaths.clear();
        let stops: Vec<(u64, Coordinate)> = self.stops.iter().map(|(k, v)| (*k, *v)).collect();
        for (a, ca) in &stops {
            for (b, cb) in &stops {
                if a == b {
                    continue;
                }
                let d = haversine_distance(ca, cb);
                if d <= max_dist_m {
                    let secs = (d / walk_speed_mps).ceil() as i64;
                    self.footpaths.entry(*a).or_default().push((*b, secs));
                }
            }
        }
    }

    /// Earliest arrival from `source` to `target` departing no earlier than
    /// `dep_time`. Returns the journey (transit + walk legs) or `None` if the
    /// target is unreachable.
    pub fn earliest_arrival(&self, source: u64, target: u64, dep_time: i64) -> Option<Journey> {
        if !self.stops.contains_key(&source) || !self.stops.contains_key(&target) {
            return None;
        }

        let mut arr: BTreeMap<u64, i64> = BTreeMap::new();
        let mut pred: BTreeMap<u64, Pred> = BTreeMap::new();
        arr.insert(source, dep_time);
        // Initial walking transfers from the source.
        self.relax_footpaths(source, dep_time, &mut arr, &mut pred);

        // Scan connections in departure order.
        let mut order: Vec<usize> = (0..self.connections.len()).collect();
        order.sort_by_key(|&i| self.connections[i].dep_time);

        for &i in &order {
            let c = self.connections[i];
            if c.dep_time < dep_time {
                continue;
            }
            let reach_dep = arr.get(&c.dep_stop).copied().unwrap_or(INF);
            if reach_dep <= c.dep_time && c.arr_time < arr.get(&c.arr_stop).copied().unwrap_or(INF) {
                arr.insert(c.arr_stop, c.arr_time);
                pred.insert(c.arr_stop, Pred::Conn(i));
                self.relax_footpaths(c.arr_stop, c.arr_time, &mut arr, &mut pred);
            }
        }

        let arrival_time = arr.get(&target).copied().unwrap_or(INF);
        if arrival_time == INF {
            return None;
        }
        self.reconstruct(source, target, arrival_time, &pred)
    }

    fn relax_footpaths(
        &self,
        from: u64,
        at: i64,
        arr: &mut BTreeMap<u64, i64>,
        pred: &mut BTreeMap<u64, Pred>,
    ) {
        if let Some(fps) = self.footpaths.get(&from) {
            for &(to, secs) in fps {
                let cand = at + secs;
                if cand < arr.get(&to).copied().unwrap_or(INF) {
                    arr.insert(to, cand);
                    pred.insert(to, Pred::Walk { from, secs });
                }
            }
        }
    }

    fn reconstruct(
        &self,
        source: u64,
        target: u64,
        arrival_time: i64,
        pred: &BTreeMap<u64, Pred>,
    ) -> Option<Journey> {
        let mut legs = Vec::new();
        let mut cur = target;
        let mut guard = 0usize;
        let limit = self.connections.len() + self.stops.len() + 2;
        while cur != source {
            guard += 1;
            if guard > limit {
                return None; // defensive: malformed predecessor chain
            }
            match pred.get(&cur)? {
                Pred::Conn(i) => {
                    let c = self.connections[*i];
                    legs.push(Leg::Transit {
                        trip: c.trip,
                        from: c.dep_stop,
                        to: c.arr_stop,
                        dep: c.dep_time,
                        arr: c.arr_time,
                    });
                    cur = c.dep_stop;
                }
                Pred::Walk { from, secs } => {
                    legs.push(Leg::Walk { from: *from, to: cur, secs: *secs });
                    cur = *from;
                }
            }
        }
        legs.reverse();
        Some(Journey { arrival_time, legs })
    }
}

/// How a stop's earliest arrival was achieved (for path reconstruction).
#[derive(Debug, Clone, Copy)]
enum Pred {
    Conn(usize),
    Walk { from: u64, secs: i64 },
}

/// One leg of a transit journey.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum Leg {
    Transit { trip: u64, from: u64, to: u64, dep: i64, arr: i64 },
    Walk { from: u64, to: u64, secs: i64 },
}

/// A computed transit journey: when it arrives and the legs to get there.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Journey {
    pub arrival_time: i64,
    pub legs: Vec<Leg>,
}

impl Journey {
    /// Number of vehicle boardings (transfers + 1, or 0 for a walk-only trip).
    pub fn boardings(&self) -> usize {
        self.legs.iter().filter(|l| matches!(l, Leg::Transit { .. })).count()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn tt() -> Timetable {
        let mut t = Timetable::default();
        for id in 1..=4 {
            t.add_stop(id, Coordinate::new(0.0, id as f64 * 0.01));
        }
        t
    }

    #[test]
    fn earliest_arrival_along_a_line_with_a_wait() {
        let mut t = tt();
        // Trip 10: stop1 @0 → stop2 @100, then (after a wait) stop2 @120 → stop3 @200.
        t.add_connection(Connection { trip: 10, dep_stop: 1, arr_stop: 2, dep_time: 0, arr_time: 100 });
        t.add_connection(Connection { trip: 10, dep_stop: 2, arr_stop: 3, dep_time: 120, arr_time: 200 });

        let j = t.earliest_arrival(1, 3, 0).expect("reachable");
        assert_eq!(j.arrival_time, 200);
        assert_eq!(j.boardings(), 2);
        assert!(matches!(j.legs[0], Leg::Transit { to: 2, .. }));
        assert!(matches!(j.legs[1], Leg::Transit { to: 3, .. }));
    }

    #[test]
    fn missing_connection_in_time_is_unreachable() {
        let mut t = tt();
        // Only a connection that departs before our departure time.
        t.add_connection(Connection { trip: 10, dep_stop: 1, arr_stop: 3, dep_time: 0, arr_time: 50 });
        assert!(t.earliest_arrival(1, 3, 100).is_none()); // missed it
        assert!(t.earliest_arrival(1, 3, 0).is_some());
    }

    #[test]
    fn picks_the_faster_of_two_trips() {
        let mut t = tt();
        // Slow trip 1→3 arriving 300; fast trip 1→3 arriving 150.
        t.add_connection(Connection { trip: 1, dep_stop: 1, arr_stop: 3, dep_time: 0, arr_time: 300 });
        t.add_connection(Connection { trip: 2, dep_stop: 1, arr_stop: 3, dep_time: 10, arr_time: 150 });
        let j = t.earliest_arrival(1, 3, 0).unwrap();
        assert_eq!(j.arrival_time, 150);
        assert!(matches!(j.legs[0], Leg::Transit { trip: 2, .. }));
    }

    #[test]
    fn transfer_via_footpath_between_trips() {
        let mut t = tt();
        // Trip A reaches stop 2 @100; trip B departs stop 4 @ 130 → stop 3 @ 200.
        // A 20s walk from stop 2 to stop 4 lets us catch B.
        t.add_connection(Connection { trip: 1, dep_stop: 1, arr_stop: 2, dep_time: 0, arr_time: 100 });
        t.add_connection(Connection { trip: 2, dep_stop: 4, arr_stop: 3, dep_time: 130, arr_time: 200 });
        t.add_footpath(2, 4, 20);

        let j = t.earliest_arrival(1, 3, 0).expect("reachable via transfer");
        assert_eq!(j.arrival_time, 200);
        // Ride, walk, ride.
        assert!(matches!(j.legs[0], Leg::Transit { to: 2, .. }));
        assert!(matches!(j.legs[1], Leg::Walk { from: 2, to: 4, secs: 20 }));
        assert!(matches!(j.legs[2], Leg::Transit { to: 3, .. }));
        assert_eq!(j.boardings(), 2);
    }

    #[test]
    fn auto_footpaths_enable_a_transfer() {
        let mut t = Timetable::default();
        // Two stops ~very close, one far.
        t.add_stop(1, Coordinate::new(0.0, 0.0));
        t.add_stop(2, Coordinate::new(0.0, 0.00001)); // ~1 m from stop 1
        t.add_stop(3, Coordinate::new(0.0, 1.0));
        t.add_connection(Connection { trip: 1, dep_stop: 2, arr_stop: 3, dep_time: 60, arr_time: 600 });
        t.auto_footpaths(50.0, 1.4); // stop1↔stop2 only

        let j = t.earliest_arrival(1, 3, 0).expect("walk to stop 2, then ride");
        assert_eq!(j.arrival_time, 600);
        assert!(matches!(j.legs[0], Leg::Walk { from: 1, to: 2, .. }));
    }
}