go-gtfs/cmd/planner/planner.go

package main

import (
	"fmt"
	"log"
	"slices"
	"sort"
	"time"

	"git.tornberg.me/go-gtfs/pkg/reader"
	"git.tornberg.me/go-gtfs/pkg/types"
)

// TripPlanner handles preprocessed transit data for efficient routing
type TripPlanner struct {
	*reader.TripData
	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
func NewTripPlanner(data *reader.TripData) *TripPlanner {
	return &TripPlanner{
		TripData: data,
		graph:    make(map[string][]Edge),
	}
}

// Preprocess builds the routing graph and precomputes routes
func (tp *TripPlanner) Preprocess() error {

	if hosjo, ok := tp.Stops["740025287"]; ok {
		trips := hosjo.GetTripsAfter(time.Now())
		for trip := range trips {
			log.Printf("Trip %s (%s):", trip.TripShortName, trip.TripHeadsign)
			for stop := range trip.GetDirectPossibleDestinations(hosjo, time.Now()) {
				log.Printf("- Stop %s at %s", stop.Stop.StopName, types.AsTime(stop.DepartureTime))
			}
		}

	}

	// Build graph with trip edges
	for tripID, trip := range tp.Trips {
		sts := trip.Stops
		sort.Slice(sts, func(i, j int) bool {
			return sts[i].StopSequence < sts[j].StopSequence
		})
		for i := 0; i < len(sts)-1; i++ {
			from := sts[i].StopId
			to := sts[i+1].StopId
			departure := sts[i].DepartureTime
			arrival := sts[i+1].DepartureTime
			timeDiff := arrival - departure
			if timeDiff > 0 {
				tp.graph[from] = append(tp.graph[from], Edge{To: to, TripID: tripID, Time: timeDiff, DepartureTime: departure})
			}
		}
	}

	// // Add transfer edges
	// for _, tr := range tp.transfers {
	// 	if tr.TransferType == 2 { // minimum transfer time
	// 		tp.graph[tr.FromStopId] = append(tp.graph[tr.FromStopId], Edge{
	// 			To:            tr.ToStopId,
	// 			TripID:        "transfer",
	// 			Time:          float64(tr.MinTransferTime),
	// 			DepartureTime: 0,
	// 		})
	// 	}
	// }

	// tp.stopTimes = nil

	return nil
}

type History struct {
	*types.StopTime
	DistanceToEnd float64
	TravelTime    types.SecondsAfterMidnight
}

func NewHistory(st *types.StopTime, distanceToEnd float64, travelTime types.SecondsAfterMidnight) *History {
	return &History{
		StopTime:      st,
		DistanceToEnd: distanceToEnd,
		TravelTime:    travelTime,
	}
}

// FindRoutes finds the best routes (up to num) between two stops starting at the given time
func (tp *TripPlanner) FindRoute(from, to string, when time.Time) (*Route, error) {

	fromStop := tp.GetStop(from)
	toStop := tp.GetStop(to)

	if fromStop == nil || toStop == nil {
		return nil, fmt.Errorf("invalid from or to stop")
	}

	possibleNextStops := make([]*types.StopTime, 0)
	var startTime *types.StopTime
	for start, stop := range fromStop.GetStopsAfter(when) {
		if stop.StopId == toStop.StopId {
			return &Route{
				Legs: []Leg{NewLeg(start, stop)},
			}, nil
		} else if from != stop.StopId {
			startTime = start
			possibleNextStops = append(possibleNextStops, stop)
		}
	}
	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
		}
	}

	return nil, fmt.Errorf("no route found")
}

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 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 {
		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
			}
		}

		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)
			}
		}
	}
	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
		}
	}
	return nil, fmt.Errorf("no route found")
}

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
	}
	legs = append(legs, NewLeg(previousStop, finalStop))
	return legs
}

func NewLeg(fromStop, toStop *types.StopTime) Leg {
	return Leg{
		From: fromStop,
		To:   toStop,
	}
}

// // findRoute implements a time-aware Dijkstra algorithm for routing
// func (tp *TripPlanner) findRoute(start, end string, when time.Time) *Route {
// 	csaPlanner := NewCSAPlanner(tp.TripData)
// 	return csaPlanner.FindRoute(start, end, when)
// }

func (tp *TripPlanner) GetRoute(routeId string) *types.Route {
	if routeId == "" {
		return nil
	}
	route, ok := tp.Routes[routeId]
	if !ok {
		return nil
	}
	return route
}

func (tp *TripPlanner) GetAgency(agencyId string) *types.Agency {
	if agencyId == "" {
		return nil
	}
	agency, ok := tp.Agencies[agencyId]
	if !ok {
		return nil
	}
	return agency
}

func (tp *TripPlanner) GetTrip(tripId string) *types.Trip {
	if tripId == "" {
		return nil
	}
	trip, ok := tp.Trips[tripId]
	if !ok {
		return nil
	}
	return trip
}

func (tp *TripPlanner) GetStop(prev string) *types.Stop {
	if prev == "" {
		return nil
	}
	stop, ok := tp.Stops[prev]
	if !ok {
		return nil
	}
	return stop
}