feat: real-time predictions and distance calculations

internal/gtfs/gtfs_manager.go

@@ -21,20 +21,23 @@ const NoRadiusLimit = -1
 
 // Manager manages the GTFS data and provides methods to access it
 type Manager struct {
-	gtfsSource       string
-	gtfsData         *gtfs.Static
-	GtfsDB           *gtfsdb.Client
-	lastUpdated      time.Time
-	isLocalFile      bool
-	realTimeTrips    []gtfs.Trip
-	realTimeVehicles []gtfs.Vehicle
-	realTimeMutex    sync.RWMutex
-	realTimeAlerts   []gtfs.Alert
-	staticMutex      sync.RWMutex // Protects gtfsData and lastUpdated
-	config           Config
-	shutdownChan     chan struct{}
-	wg               sync.WaitGroup
-	shutdownOnce     sync.Once
+	gtfsSource                     string
+	gtfsData                       *gtfs.Static
+	GtfsDB                         *gtfsdb.Client
+	lastUpdated                    time.Time
+	isLocalFile                    bool
+	realTimeTrips                  []gtfs.Trip
+	realTimeVehicles               []gtfs.Vehicle
+	realTimeMutex                  sync.RWMutex
+	realTimeAlerts                 []gtfs.Alert
+	realTimeTripLookup             map[string]int
+	realTimeVehicleLookupByTrip    map[string]int
+	realTimeVehicleLookupByVehicle map[string]int
+	staticMutex                    sync.RWMutex // Protects gtfsData and lastUpdated
+	config                         Config
+	shutdownChan                   chan struct{}
+	wg                             sync.WaitGroup
+	shutdownOnce                   sync.Once
 }
 
 // InitGTFSManager initializes the Manager with the GTFS data from the given source
@@ -48,10 +51,13 @@ func InitGTFSManager(config Config) (*Manager, error) {
 	}
 
 	manager := &Manager{
-		gtfsSource:   config.GtfsURL,
-		isLocalFile:  isLocalFile,
-		config:       config,
-		shutdownChan: make(chan struct{}),
+		gtfsSource:                     config.GtfsURL,
+		isLocalFile:                    isLocalFile,
+		config:                         config,
+		shutdownChan:                   make(chan struct{}),
+		realTimeTripLookup:             make(map[string]int),
+		realTimeVehicleLookupByTrip:    make(map[string]int),
+		realTimeVehicleLookupByVehicle: make(map[string]int),
 	}
 	manager.setStaticGTFS(staticData)
 
@@ -308,10 +314,8 @@ func (manager *Manager) GetVehicleByID(vehicleID string) (*gtfs.Vehicle, error)
 	manager.realTimeMutex.RLock()
 	defer manager.realTimeMutex.RUnlock()
 
-	for _, v := range manager.realTimeVehicles {
-		if v.ID.ID == vehicleID {
-			return &v, nil
-		}
+	if index, exists := manager.realTimeVehicleLookupByVehicle[vehicleID]; exists {
+		return &manager.realTimeVehicles[index], nil
 	}
 
 	return nil, fmt.Errorf("vehicle with ID %s not found", vehicleID)
@@ -330,6 +334,22 @@ func (manager *Manager) GetTripUpdatesForTrip(tripID string) []gtfs.Trip {
 	return updates
 }
 
+func (manager *Manager) GetVehicleLastUpdateTime(vehicle *gtfs.Vehicle) int64 {
+	if vehicle == nil || vehicle.Timestamp == nil {
+		return 0
+	}
+	return vehicle.Timestamp.UnixMilli()
+}
+
+func (manager *Manager) GetTripUpdateByID(tripID string) (*gtfs.Trip, error) {
+	manager.realTimeMutex.RLock()
+	defer manager.realTimeMutex.RUnlock()
+	if index, exists := manager.realTimeTripLookup[tripID]; exists {
+		return &manager.realTimeTrips[index], nil
+	}
+	return nil, fmt.Errorf("trip with ID %s not found", tripID)
+}
+
 func (manager *Manager) PrintStatistics() {
 	manager.staticMutex.RLock()
 	defer manager.staticMutex.RUnlock()

internal/gtfs/gtfs_manager_mock.go

@@ -43,6 +43,8 @@ func (m *Manager) MockAddVehicle(vehicleID, tripID, routeID string) {
 			},
 		},
 	})
+
+	m.realTimeVehicleLookupByVehicle[vehicleID] = len(m.realTimeVehicles) - 1
 }
 
 func (m *Manager) MockAddTrip(tripID, agencyID, routeID string) {

internal/gtfs/realtime.go

@@ -166,16 +166,61 @@ func (manager *Manager) updateGTFSRealtime(ctx context.Context, config Config) {
 
 	if tripData != nil && tripErr == nil {
 		manager.realTimeTrips = tripData.Trips
+		rebuildRealTimeTripLookup(manager)
+
 	}
 	if vehicleData != nil && vehicleErr == nil {
 		manager.realTimeVehicles = vehicleData.Vehicles
+		rebuildRealTimeVehicleLookupByTrip(manager)
+		rebuildRealTimeVehicleLookupByVehicle(manager)
 	}
 
 	if alertData != nil && alertErr == nil {
 		manager.realTimeAlerts = alertData.Alerts
 	}
 }
 
+func rebuildRealTimeTripLookup(manager *Manager) {
+	if manager.realTimeTripLookup == nil {
+		manager.realTimeTripLookup = make(map[string]int)
+	} else {
+		for k := range manager.realTimeTripLookup {
+			delete(manager.realTimeTripLookup, k)
+		}
+	}
+	for i, trip := range manager.realTimeTrips {
+		manager.realTimeTripLookup[trip.ID.ID] = i
+	}
+}
+
+func rebuildRealTimeVehicleLookupByTrip(manager *Manager) {
+	if manager.realTimeVehicleLookupByTrip == nil {
+		manager.realTimeVehicleLookupByTrip = make(map[string]int)
+	} else {
+		for k := range manager.realTimeVehicleLookupByTrip {
+			delete(manager.realTimeVehicleLookupByTrip, k)
+		}
+	}
+	for i, vehicle := range manager.realTimeVehicles {
+		if vehicle.Trip != nil && vehicle.Trip.ID.ID != "" {
+			manager.realTimeVehicleLookupByTrip[vehicle.Trip.ID.ID] = i
+		}
+	}
+}
+
+func rebuildRealTimeVehicleLookupByVehicle(manager *Manager) {
+	if manager.realTimeVehicleLookupByVehicle == nil {
+		manager.realTimeVehicleLookupByVehicle = make(map[string]int)
+	} else {
+		for k := range manager.realTimeVehicleLookupByVehicle {
+			delete(manager.realTimeVehicleLookupByVehicle, k)
+		}
+	}
+	for i, vehicle := range manager.realTimeVehicles {
+		manager.realTimeVehicleLookupByVehicle[vehicle.ID.ID] = i
+	}
+}
+
 func (manager *Manager) updateGTFSRealtimePeriodically(config Config) {
 	defer manager.wg.Done()
 

internal/models/trip_details.go

@@ -33,31 +33,33 @@ func NewEmptyTripDetails() *TripDetails {
 }
 
 type TripStatusForTripDetails struct {
-	ServiceDate                int64    `json:"serviceDate"`
-	ActiveTripID               string   `json:"activeTripId"`
-	Phase                      string   `json:"phase"`
-	Status                     string   `json:"status"`
-	Predicted                  bool     `json:"predicted"`
-	VehicleID                  string   `json:"vehicleId"`
-	Position                   Location `json:"position"`
-	LastKnownLocation          Location `json:"lastKnownLocation"`
-	Orientation                float64  `json:"orientation"`
-	LastKnownOrientation       float64  `json:"lastKnownOrientation"`
-	ScheduleDeviation          int      `json:"scheduleDeviation"`
-	DistanceAlongTrip          float64  `json:"distanceAlongTrip"`
-	ScheduledDistanceAlongTrip float64  `json:"scheduledDistanceAlongTrip"`
-	TotalDistanceAlongTrip     float64  `json:"totalDistanceAlongTrip"`
-	LastKnownDistanceAlongTrip float64  `json:"lastKnownDistanceAlongTrip"`
-	LastUpdateTime             int64    `json:"lastUpdateTime"`
-	LastLocationUpdateTime     int64    `json:"lastLocationUpdateTime"`
-	BlockTripSequence          int      `json:"blockTripSequence"`
-	ClosestStop                string   `json:"closestStop"`
-	ClosestStopTimeOffset      int      `json:"closestStopTimeOffset"`
-	NextStop                   string   `json:"nextStop"`
-	NextStopTimeOffset         int      `json:"nextStopTimeOffset"`
-	OccupancyStatus            string   `json:"occupancyStatus"`
-	OccupancyCount             int      `json:"occupancyCount"`
-	OccupancyCapacity          int      `json:"occupancyCapacity"`
-	SituationIDs               []string `json:"situationIds"`
-	Scheduled                  bool     `json:"scheduled"`
+	ActiveTripID               string     `json:"activeTripId"`
+	BlockTripSequence          int        `json:"blockTripSequence"`
+	ClosestStop                string     `json:"closestStop"`
+	ClosestStopTimeOffset      int        `json:"closestStopTimeOffset"`
+	DistanceAlongTrip          float64    `json:"distanceAlongTrip"`
+	Frequency                  *Frequency `json:"frequency,omitempty"`
+	LastKnownDistanceAlongTrip float64    `json:"lastKnownDistanceAlongTrip"`
+	LastKnownLocation          Location   `json:"lastKnownLocation"`
+	LastKnownOrientation       float64    `json:"lastKnownOrientation"`
+	LastLocationUpdateTime     int64      `json:"lastLocationUpdateTime"`
+	LastUpdateTime             int64      `json:"lastUpdateTime"`
+	NextStop                   string     `json:"nextStop"`
+	NextStopTimeOffset         int        `json:"nextStopTimeOffset"`
+	OccupancyCapacity          int        `json:"occupancyCapacity"`
+	OccupancyCount             int        `json:"occupancyCount"`
+	OccupancyStatus            string     `json:"occupancyStatus"`
+	Orientation                float64    `json:"orientation"`
+	Phase                      string     `json:"phase"`
+	Position                   Location   `json:"position"`
+	Predicted                  bool       `json:"predicted"`
+	ScheduleDeviation          int        `json:"scheduleDeviation"`
+	ScheduledDistanceAlongTrip float64    `json:"scheduledDistanceAlongTrip"`
+	ServiceDate                int64      `json:"serviceDate"`
+	SituationIDs               []string   `json:"situationIds"`
+	Status                     string     `json:"status"`
+	TotalDistanceAlongTrip     float64    `json:"totalDistanceAlongTrip"`
+	VehicleFeatures            []string   `json:"vehicleFeatures,omitempty"`
+	VehicleID                  string     `json:"vehicleId"`
+	Scheduled                  bool       `json:"scheduled"`
 }