some custom stuff

pkg/types/stop.go

@@ -1,17 +1,18 @@
 package types
 
 import (
+	"iter"
 	"time"
 )
 
 type Stop struct {
-	trips        map[string]*Trip
-	StopID       string      `json:"stop_id" csv:"stop_id"`
-	StopName     string      `json:"stop_name" csv:"stop_name"`
-	StopLat      float64     `json:"stop_lat" csv:"stop_lat"`
-	StopLon      float64     `json:"stop_lon" csv:"stop_lon"`
-	LocationType int         `json:"location_type" csv:"location_type"`
-	Transfers    []*Transfer `json:"transfers" csv:"transfers"`
+	Trips        map[string]*Trip `json:"-" csv:"-"`
+	StopId       string           `json:"stop_id" csv:"stop_id"`
+	StopName     string           `json:"stop_name" csv:"stop_name"`
+	StopLat      float64          `json:"stop_lat" csv:"stop_lat"`
+	StopLon      float64          `json:"stop_lon" csv:"stop_lon"`
+	LocationType int              `json:"location_type" csv:"location_type"`
+	Transfers    []*Transfer      `json:"-" csv:"transfers"`
 }
 
 func (s *Stop) AddTransfer(transfer *Transfer) {
@@ -22,21 +23,71 @@ func (s *Stop) AddTransfer(transfer *Transfer) {
 }
 
 func (s *Stop) AddTrip(trip *Trip) {
-	if s.trips == nil {
-		s.trips = make(map[string]*Trip)
-	}
-	s.trips[trip.TripID] = trip
+
+	s.Trips[trip.TripId] = trip
 }
 
-func (s *Stop) GetTripsAfter(time time.Time) []*Trip {
-	var trips []*Trip
-	for _, trip := range s.trips {
-		for _, stop := range trip.Stops {
-			if stop.StopID == s.StopID && stop.DepartsAfter(time) {
-				trips = append(trips, trip)
-				break
+type TripWithDepartureTime struct {
+	*Trip
+	DepartureTime SecondsAfterMidnight
+}
+
+func (s *Stop) GetTripsAfter(when time.Time) iter.Seq[*TripWithDepartureTime] {
+	startAfterMidnight := AsSecondsAfterMidnight(when)
+	return func(yield func(*TripWithDepartureTime) bool) {
+		for _, trip := range s.Trips {
+
+			for _, stop := range trip.Stops {
+				if stop.StopId == s.StopId && stop.ArrivalTime >= startAfterMidnight {
+
+					if !yield(&TripWithDepartureTime{Trip: trip, DepartureTime: stop.DepartureTime}) {
+						return
+					}
+					break
+				}
+			}
+
+		}
+	}
+}
+
+func (s *Stop) GetUpcomingStops(start *StopTime) iter.Seq[*StopTime] {
+	return func(yield func(*StopTime) bool) {
+		found := false
+		for _, trip := range s.Trips {
+			for _, stop := range trip.Stops {
+				if !found {
+					if stop.StopId == start.StopId && stop.DepartureTime >= start.ArrivalTime {
+						found = true
+					}
+				} else {
+					if !yield(stop) {
+						return
+					}
+				}
 			}
 		}
 	}
-	return trips
+}
+
+func (s *Stop) GetStopsAfter(when time.Time) iter.Seq2[*StopTime, *StopTime] {
+	startAfterMidnight := AsSecondsAfterMidnight(when)
+	return func(yield func(start, stop *StopTime) bool) {
+		for _, trip := range s.Trips {
+			found := false
+			var start *StopTime
+			for _, stop := range trip.Stops {
+				if stop.StopId == s.StopId && stop.ArrivalTime >= startAfterMidnight {
+					found = true
+					start = stop
+				}
+				if found {
+					if !yield(start, stop) {
+						return
+					}
+				}
+			}
+		}
+
+	}
 }