go-gtfs/cmd/planner/planner.go

package main

import (
	"fmt"
	"log"
	"slices"
	"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")
	}
	routes := make([]*Route, 0)
	usedRouteIds := make(map[string]struct{})

	for trip := range fromStop.GetTripsAfter(when) {
		if _, used := usedRouteIds[trip.RouteId]; used {
			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)

			}
		}

	}

	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 err == nil && route != nil {
				routes = append(routes, route)
			}
		}
	}
	// 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
		}
		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
			}
		}

		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 {
	trip, ok := toStop.Stop.Trips[toStop.TripId]
	if !ok {
		log.Printf("trip %s not found for stop %s", toStop.TripId, toStop.Stop.StopName)
		return Leg{
			From: fromStop,
			To:   toStop,
		}
	}
	return Leg{
		From: fromStop,
		To:   toStop,
		Trip: &JSONTrip{
			TripId:        trip.TripId,
			RouteId:       trip.RouteID,
			AgencyName:    trip.Agency.AgencyName,
			TripHeadsign:  trip.TripHeadsign,
			TripShortName: trip.TripShortName,
		},
	}
}

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