fix

cmd/planner/planner.go

@@ -1,9 +1,9 @@
 package main
 
 import (
+	"container/heap"
 	"fmt"
 	"log"
-	"slices"
 	"time"
 
 	"git.tornberg.me/go-gtfs/pkg/reader"
@@ -16,24 +16,12 @@ type TripPlanner struct {
 	graph map[string][]Edge
 }
 
-type StopWithPossibleConnections struct {
-	*types.Stop
-	PossibleConnections []Connection
-}
-
-type Connection struct {
-	*types.Stop
-	Distance float64
-	Time     time.Duration
-}
-
 const (
 	transferPenalty     = 90 * time.Minute
 	maxTransfers        = 4
 	maxWaitBetweenTrips = 1 * time.Hour
 	//trajectoryAngleTolerance = 220.0
 	maxTravelDuration = 12 * time.Hour
-	maxDetourFactor   = 2
 )
 
 // NewTripPlanner creates a new trip planner instance
@@ -88,23 +76,41 @@ func (tp *TripPlanner) Preprocess() error {
 	// 	}
 	// }
 
-	// tp.stopTimes = nil
-
 	return nil
 }
 
-type History struct {
-	*types.StopTime
-	DistanceToEnd float64
-	TravelTime    types.SecondsAfterMidnight
+type searchNode struct {
+	stopTime  *types.StopTime
+	index     int
+	g         float64
+	f         float64
+	transfers int
+	parent    *searchNode
 }
 
-func NewHistory(st *types.StopTime, distanceToEnd float64, travelTime types.SecondsAfterMidnight) History {
-	return History{
-		StopTime:      st,
-		DistanceToEnd: distanceToEnd,
-		TravelTime:    travelTime,
+func nodeKey(st *types.StopTime) string {
+	return fmt.Sprintf("%s:%d", st.TripId, st.StopSequence)
+}
+
+func heuristicSeconds(from, to *types.Stop) float64 {
+	if from == nil || to == nil {
+		return 0
 	}
+	const avgSpeedKmh = 60.0
+	distanceKm := from.HaversineDistance(to)
+	return (distanceKm / avgSpeedKmh) * 3600
+}
+
+func tripStopIndex(trip *types.Trip, st *types.StopTime) int {
+	if trip == nil {
+		return -1
+	}
+	for i, candidate := range trip.Stops {
+		if candidate == st {
+			return i
+		}
+	}
+	return -1
 }
 
 // FindRoutes finds the best routes (up to num) between two stops starting at the given time
@@ -116,160 +122,163 @@ func (tp *TripPlanner) FindRoute(from, to string, when time.Time) ([]*Route, err
 	if fromStop == nil || toStop == nil {
 		return nil, fmt.Errorf("invalid from or to stop")
 	}
-	routes := make([]*Route, 0)
-	usedRouteIds := make(map[string]struct{})
+	open := priorityQueue{}
+	heap.Init(&open)
+	bestCost := make(map[string]float64)
+	startSeconds := types.AsSecondsAfterMidnight(when)
+	initCount := 0
 
-	for trip := range fromStop.GetTripsAfter(when) {
-		if _, used := usedRouteIds[trip.RouteId]; used {
+	for tripWithDeparture := range fromStop.GetTripsAfter(when) {
+		trip := tripWithDeparture.Trip
+		startStopTime, ok := trip.Has(fromStop)
+		if !ok {
 			continue
 		}
-		for i := len(trip.Stops) - 1; i >= 0; i-- {
-			stop := trip.Stops[i]
-			if stop.StopId == toStop.StopId {
-				usedRouteIds[trip.RouteId] = struct{}{}
-				routes = append(routes, &Route{
-					Legs: []Leg{NewLeg(trip.Stops[0], trip.Stops[i])},
-				})
-				break
-			} else if stop.StopId == fromStop.StopId {
-				break
-			} else if stop.PickupType == 0 {
-				distance := stop.Stop.HaversineDistance(toStop)
-
-			}
+		idx := tripStopIndex(trip, startStopTime)
+		if idx < 0 {
+			continue
 		}
-
+		wait := float64(startStopTime.DepartureTime - startSeconds)
+		if wait < 0 {
+			wait = 0
+		}
+		node := &searchNode{
+			stopTime:  startStopTime,
+			index:     idx,
+			g:         wait,
+			transfers: 0,
+		}
+		node.f = node.g + heuristicSeconds(startStopTime.Stop, toStop)
+		heap.Push(&open, &pqItem{node: node, priority: node.f})
+		bestCost[nodeKey(startStopTime)] = node.g
+		initCount++
 	}
 
-	for start, stop := range fromStop.GetStopsAfter(when) {
-		if stop.StopId == toStop.StopId {
-			routes = append(routes, &Route{
-				Legs: []Leg{NewLeg(start, stop)},
-			})
-		} else if from != stop.StopId {
-			//		startTime = start
-			route, err := tp.findRoute(*start, toStop, NewHistory(start, start.Stop.HaversineDistance(toStop), 0), NewHistory(stop, stop.Stop.HaversineDistance(toStop), stop.ArrivalTime-start.DepartureTime))
+	if initCount == 0 {
+		return nil, fmt.Errorf("no departures from %s after %s", from, when.Format(time.RFC3339))
+	}
 
-			if err == nil && route != nil {
+	routes := make([]*Route, 0, 3)
+
+	for open.Len() > 0 && len(routes) < 3 {
+		current := heap.Pop(&open).(*pqItem).node
+		if current.stopTime.StopId == toStop.StopId {
+			route := buildRouteFromNode(current)
+			if route != nil {
 				routes = append(routes, route)
 			}
+			continue
 		}
-	}
-	// slices.SortFunc(possibleNextStops, byDistanceTo(*toStop))
-	// for _, nextStop := range possibleNextStops {
-	// 	route, err := tp.findRoute(*nextStop, toStop, *NewHistory(startTime, startTime.Stop.HaversineDistance(toStop), types.AsSecondsAfterMidnight(when)), *NewHistory(nextStop, nextStop.Stop.HaversineDistance(toStop), nextStop.ArrivalTime-startTime.ArrivalTime))
-	// 	if err == nil && route != nil {
-	// 		return route, nil
-	// 	}
-	// }
-	slices.SortFunc(routes, func(a, b *Route) int {
-		transfersA := len(a.Legs) - 1
-		transfersB := len(b.Legs) - 1
-		if transfersA != transfersB {
-			return transfersA - transfersB
+		if current.g > maxTravelDuration.Seconds() {
+			continue
 		}
-		return a.Duration() - b.Duration() - (transfersA-transfersB)*int(transferPenalty.Seconds())
-	})
-	return routes[:min(len(routes), 10)], nil
-}
-
-func byDistanceTo(end types.Stop) func(a, b *types.StopTime) int {
-	return func(a, b *types.StopTime) int {
-		distanceA := haversine(a.Stop.StopLat, a.Stop.StopLon, end.StopLat, end.StopLon)
-		distanceB := haversine(b.Stop.StopLat, b.Stop.StopLon, end.StopLat, end.StopLon)
-		return (int(distanceA) - int(distanceB)) // + (int(b.ArrivalTime - a.ArrivalTime))
-	}
-}
-
-func isInCorrectDirection(from, possible, end *types.Stop) bool {
-	if from.StopId == end.StopId || possible.StopId == end.StopId {
-		return true
-	}
-	if from.StopId == possible.StopId {
-		return false
-	}
-	startToEndLat := end.StopLat - from.StopLat
-	startToEndLon := end.StopLon - from.StopLon
-	startToPossibleLat := possible.StopLat - from.StopLat
-	startToPossibleLon := possible.StopLon - from.StopLon
-	dotProduct := startToEndLat*startToPossibleLat + startToEndLon*startToPossibleLon
-	return dotProduct > -0.4 && dotProduct < 0.4
-}
-
-func shouldTryStop(end *types.Stop, visited ...History) func(possible *types.StopTime) bool {
-	lastDistance := visited[len(visited)-1].Stop.HaversineDistance(end)
-	return func(possible *types.StopTime) bool {
-		if end.StopId == possible.StopId {
-			return true
-		}
-		if possible.DepartureTime > visited[len(visited)-1].DepartureTime+types.SecondsAfterMidnight(maxWaitBetweenTrips.Seconds()) {
-			return false
-		}
-		if possible.DropOffType == 1 {
-			return false
-		}
-		// if !isInCorrectDirection(visited[len(visited)-1].Stop, possible.Stop, end) {
-		// 	return false
-		// }
-		distance := possible.Stop.HaversineDistance(end)
-
-		for _, v := range visited {
-			if v.DistanceToEnd <= distance*1.2 {
-				return false
-			}
-			if v.TripId == possible.TripId || v.StopId == possible.StopId {
-				return false
+		tpp := tp.Trips[current.stopTime.TripId]
+		if tpp != nil {
+			for i := current.index + 1; i < len(tpp.Stops); i++ {
+				next := tpp.Stops[i]
+				if next.DropOffType == 1 {
+					continue
+				}
+				travel := next.ArrivalTime - current.stopTime.DepartureTime
+				if travel <= 0 {
+					continue
+				}
+				newCost := current.g + float64(travel)
+				if newCost > maxTravelDuration.Seconds() {
+					continue
+				}
+				key := nodeKey(next)
+				if prev, ok := bestCost[key]; ok && newCost >= prev {
+					continue
+				}
+				nextNode := &searchNode{
+					stopTime:  next,
+					index:     i,
+					g:         newCost,
+					transfers: current.transfers,
+					parent:    current,
+				}
+				nextNode.f = newCost + heuristicSeconds(next.Stop, toStop)
+				bestCost[key] = newCost
+				heap.Push(&open, &pqItem{node: nextNode, priority: nextNode.f})
 			}
 		}
-
-		return distance <= lastDistance*1.2
-	}
-}
-
-func (tp *TripPlanner) findRoute(start types.StopTime, end *types.Stop, changes ...History) (*Route, error) {
-	if len(changes) >= maxTransfers {
-		return nil, fmt.Errorf("max transfers reached")
-	}
-	isOk := shouldTryStop(end, changes...)
-	possibleNextStops := make([]*types.StopTime, 0)
-	for stop := range start.Stop.GetUpcomingStops(&start) {
-		if stop.StopId == end.StopId {
-			return &Route{
-				Legs: CreateLegs(changes, stop),
-			}, nil
-		} else {
-			if isOk(stop) {
-				possibleNextStops = append(possibleNextStops, stop)
+		if current.transfers >= maxTransfers {
+			continue
+		}
+		for _, otherTrip := range current.stopTime.Stop.Trips {
+			if otherTrip.TripId == current.stopTime.TripId {
+				continue
 			}
+			transferStopTime, ok := otherTrip.Has(current.stopTime.Stop)
+			if !ok {
+				continue
+			}
+			if transferStopTime.DepartureTime <= current.stopTime.ArrivalTime {
+				continue
+			}
+			wait := transferStopTime.DepartureTime - current.stopTime.ArrivalTime
+			if wait > types.SecondsAfterMidnight(maxWaitBetweenTrips.Seconds()) {
+				continue
+			}
+			newCost := current.g + float64(wait) + transferPenalty.Seconds()
+			if newCost > maxTravelDuration.Seconds() {
+				continue
+			}
+			idx := tripStopIndex(otherTrip, transferStopTime)
+			if idx < 0 {
+				continue
+			}
+			key := nodeKey(transferStopTime)
+			if prev, ok := bestCost[key]; ok && newCost >= prev {
+				continue
+			}
+			nextNode := &searchNode{
+				stopTime:  transferStopTime,
+				index:     idx,
+				g:         newCost,
+				transfers: current.transfers + 1,
+				parent:    current,
+			}
+			nextNode.f = newCost + heuristicSeconds(transferStopTime.Stop, toStop)
+			bestCost[key] = newCost
+			heap.Push(&open, &pqItem{node: nextNode, priority: nextNode.f})
 		}
 	}
-	slices.SortFunc(possibleNextStops, byDistanceTo(*end))
 
-	tries := 15
-	for _, nextStop := range possibleNextStops {
-		route, err := tp.findRoute(*nextStop, end, append(changes, NewHistory(nextStop, nextStop.Stop.HaversineDistance(end), nextStop.ArrivalTime-start.ArrivalTime))...)
-		if err == nil && route != nil {
-			return route, nil
-		}
-		tries--
-		if tries <= 0 {
-			break
-		}
+	if len(routes) == 0 {
+		return nil, fmt.Errorf("no route found")
 	}
-	return nil, fmt.Errorf("no route found")
+
+	return routes, nil
 }
 
-func CreateLegs(stops []History, finalStop *types.StopTime) []Leg {
-	legs := make([]Leg, 0, len(stops)+1)
-	var previousStop *types.StopTime
-	for _, stop := range stops {
-		if previousStop != nil {
-			legs = append(legs, NewLeg(previousStop, stop.StopTime))
-		}
-		previousStop = stop.StopTime
+func buildRouteFromNode(goal *searchNode) *Route {
+	if goal == nil {
+		return nil
 	}
-	legs = append(legs, NewLeg(previousStop, finalStop))
-	return legs
+	nodes := make([]*searchNode, 0)
+	for current := goal; current != nil; current = current.parent {
+		nodes = append(nodes, current)
+	}
+	for i, j := 0, len(nodes)-1; i < j; i, j = i+1, j-1 {
+		nodes[i], nodes[j] = nodes[j], nodes[i]
+	}
+	if len(nodes) < 2 {
+		return nil
+	}
+	legs := make([]Leg, 0)
+	runStart := nodes[0]
+	for i := 1; i < len(nodes); i++ {
+		prev := nodes[i-1]
+		curr := nodes[i]
+		if curr.stopTime.TripId != prev.stopTime.TripId {
+			legs = append(legs, NewLeg(runStart.stopTime, prev.stopTime))
+			runStart = curr
+		}
+	}
+	legs = append(legs, NewLeg(runStart.stopTime, nodes[len(nodes)-1].stopTime))
+	return &Route{Legs: legs}
 }
 
 func NewLeg(fromStop, toStop *types.StopTime) Leg {
@@ -281,13 +290,17 @@ func NewLeg(fromStop, toStop *types.StopTime) Leg {
 			To:   toStop,
 		}
 	}
+	agencyName := ""
+	if trip.Agency != nil {
+		agencyName = trip.Agency.AgencyName
+	}
 	return Leg{
 		From: fromStop,
 		To:   toStop,
 		Trip: &JSONTrip{
 			TripId:        trip.TripId,
-			RouteId:       trip.RouteID,
-			AgencyName:    trip.Agency.AgencyName,
+			RouteId:       trip.RouteId,
+			AgencyName:    agencyName,
 			TripHeadsign:  trip.TripHeadsign,
 			TripShortName: trip.TripShortName,
 		},