go-gtfs/cmd/planner/main.go

package main

import (
	"container/heap"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"sort"
	"strconv"
	"sync"
	"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 {
	stops         map[string]*types.Stop
	trips         map[string]*types.Trip
	routes        map[string]*types.Route
	agencies      map[string]*types.Agency
	transfers     []types.Transfer
	stopTimes     map[string][]types.StopTime
	graph         map[string][]Edge
	weightedStops map[string][]Connection
}

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() *TripPlanner {
	return &TripPlanner{
		stops:         make(map[string]*types.Stop),
		trips:         make(map[string]*types.Trip),
		routes:        make(map[string]*types.Route),
		agencies:      make(map[string]*types.Agency),
		stopTimes:     make(map[string][]types.StopTime),
		graph:         make(map[string][]Edge),
		weightedStops: make(map[string][]Connection),
	}
}

// LoadData loads all GTFS data
func (tp *TripPlanner) LoadData(dataDir string) error {
	files := []string{"agency", "routes", "stops", "trips", "stop_times", "transfers"}
	for _, file := range files {

		f, err := os.Open(dataDir + "/" + file + ".txt")
		if err != nil {
			log.Fatalf("failed to open %s: %v", file, err)
		}

		switch file {
		case "agency":
			err = reader.ParseAgencies(f, func(a types.Agency) {
				tp.agencies[a.AgencyID] = &a
			})
		case "routes":
			err = reader.ParseRoutes(f, func(r types.Route) {
				tp.routes[r.RouteID] = &r
				if ag, ok := tp.agencies[r.AgencyID]; ok {
					r.Agency = ag
					ag.AddRoute(&r)
				}
			})
		case "stops":
			err = reader.ParseStops(f, func(s types.Stop) {
				tp.stops[s.StopID] = &s
			})
		case "trips":
			err = reader.ParseTrips(f, func(t types.Trip) {
				trip := t
				if route, ok := tp.routes[trip.RouteID]; ok {
					trip.SetRoute(route)
					route.AddTrip(&trip)
				} else {
					log.Printf("route %s not found", trip.RouteID)
				}
				if agency, ok := tp.agencies[trip.AgencyID]; ok {
					trip.Agency = agency
				} else {
					log.Printf("agency %s not found", trip.AgencyID)
				}
				tp.trips[trip.TripID] = &trip
			})
		case "stop_times":
			err = reader.ParseStopTimes(f, func(st types.StopTime) {
				stop, ok := tp.stops[st.StopID]
				if ok {
					st.SetStop(stop)
				} else {
					log.Printf("stop %s not found", st.StopID)

				}
				trp, ok := tp.trips[st.TripID]
				if !ok {
					log.Printf("trip %s not found", st.TripID)

				} else {
					stop.AddTrip(trp)
					trp.AddStopTime(&st)
				}

			})
		case "transfers":
			err = reader.ParseTransfers(f, func(tr types.Transfer) {
				tp.transfers = append(tp.transfers, tr)
				stop, ok := tp.stops[tr.FromStopID]
				if ok {
					stop.AddTransfer(&tr)
				} else {
					log.Printf("stop %s not found for transfer", tr.FromStopID)
				}

			})
		}
		if err != nil {
			log.Printf("failed to parse %s: %v", file, err)
		}
		f.Close()

	}
	return nil
}

// 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, stop.DepartureTime)
			}
		}

	}

	// Build graph with trip edges
	for tripID, sts := range tp.stopTimes {
		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 := parseTime(sts[i].DepartureTime)
			arrival := parseTime(sts[i+1].ArrivalTime)
			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
}

// FindRoute finds the best route between two stops starting at the given time
func (tp *TripPlanner) FindRoute(from, to string, when time.Time) *Route {
	routes := tp.FindRoutes(from, to, when, 1)
	if len(routes) > 0 {
		return routes[0]
	}
	return nil
}

// FindRoutes finds the best routes (up to num) between two stops starting at the given time
func (tp *TripPlanner) FindRoutes(from, to string, when time.Time, num int) []*Route {
	var allRoutes []*Route
	seen := make(map[string]bool) // to avoid duplicates based on departure and arrival times
	for i := 0; i < num*20 && len(allRoutes) < num; i++ {
		route := tp.findRoute(from, to, when.Add(time.Duration(i*5)*time.Minute))
		if route != nil {
			key := fmt.Sprintf("%d-%d", route.Legs[0].DepartureTime.Unix(), route.Legs[len(route.Legs)-1].ArrivalTime.Unix())
			if !seen[key] {
				seen[key] = true
				allRoutes = append(allRoutes, route)
				log.Printf("Found route %d: departure %v, arrival %v, duration %v", len(allRoutes), route.Legs[0].DepartureTime, route.Legs[len(route.Legs)-1].ArrivalTime, route.Duration())
			}
		}
	}
	log.Printf("Total routes found: %d", len(allRoutes))

	return allRoutes
}

// findRoute implements a time-aware Dijkstra algorithm for routing
func (tp *TripPlanner) findRoute(start, end string, when time.Time) *Route {
	startStop := tp.GetStop(start)
	if startStop == nil {
		return nil
	}
	goalStop := tp.GetStop(end)
	if goalStop == nil {
		return nil
	}
	maxAllowedDistance := haversine(startStop.StopLat, startStop.StopLon, goalStop.StopLat, goalStop.StopLon) * maxDetourFactor

	arrival := make(map[string]time.Time)
	cost := make(map[string]time.Time)
	prev := make(map[string]PathInfo)
	pq := &priorityQueue{}
	heap.Init(pq)

	arrival[start] = when
	cost[start] = when
	prev[start] = PathInfo{Prev: "", TripID: "", DepartureTime: when, Transfers: 0, LastTrip: "", WaitDuration: 0}
	heap.Push(pq, &pqItem{Stop: start, Cost: when})

	for pq.Len() > 0 {
		item := heap.Pop(pq).(*pqItem)
		current := item.Stop
		if storedCost, ok := cost[current]; !ok || item.Cost.After(storedCost) {
			continue
		}

		currentArrival, ok := arrival[current]
		if !ok || currentArrival.IsZero() {

			continue
		}

		if current == end {
			stopsPath := []string{}
			current := end
			visited := make(map[string]bool)
			for current != "" {
				if visited[current] {
					break
				}
				visited[current] = true
				stopsPath = append([]string{current}, stopsPath...)
				if current == start {
					break
				}
				info, ok := prev[current]
				if !ok || info.Prev == "" {
					break
				}
				current = info.Prev
			}
			legs := []Leg{}
			var currentLeg *Leg
			for i := 0; i < len(stopsPath)-1; i++ {
				from := stopsPath[i]
				to := stopsPath[i+1]
				tripID := prev[to].TripID
				if tripID != "transfer" {
					if currentLeg == nil || currentLeg.TripID != tripID {
						if currentLeg != nil {
							legs = append(legs, *currentLeg)
						}
						trip := tp.GetTrip(tripID)
						route := tp.GetRoute(trip.RouteID)
						currentLeg = &Leg{
							TripID:        tripID,
							From:          from,
							FromStop:      tp.GetStop(from),
							Trip:          trip,
							Agency:        tp.GetAgency(route.AgencyID),
							Route:         route,
							Stops:         []string{from},
							DepartureTime: prev[to].DepartureTime,
						}
					}
					currentLeg.To = to
					currentLeg.ToStop = tp.GetStop(to)
					currentLeg.Stops = append(currentLeg.Stops, to)
					currentLeg.ArrivalTime = arrival[to]
				}
			}
			if currentLeg != nil {
				currentLeg.To = stopsPath[len(stopsPath)-1]
				currentLeg.ToStop = tp.GetStop(currentLeg.To)
				currentLeg.ArrivalTime = arrival[stopsPath[len(stopsPath)-1]]
				legs = append(legs, *currentLeg)
			}

			return &Route{Legs: legs}
		}

		currentStop := tp.GetStop(current)
		if currentStop == nil {
			continue
		}
		//currentBearing := bearing(currentStop.StopLat, currentStop.StopLon, goalStop.StopLat, goalStop.StopLon)

		for _, edge := range tp.graph[current] {
			if edge.To == current {
				continue
			}
			if info, ok := prev[current]; ok && info.Prev == edge.To && info.TripID == edge.TripID {
				continue
			}
			nextStop := tp.GetStop(edge.To)
			if nextStop == nil {
				continue
			}

			distanceToGoal := haversine(nextStop.StopLat, nextStop.StopLon, goalStop.StopLat, goalStop.StopLon)
			if distanceToGoal > maxAllowedDistance {
				continue
			}

			// if edge.TripID != "transfer" {
			// 	edgeBearing := bearing(currentStop.StopLat, currentStop.StopLon, nextStop.StopLat, nextStop.StopLon)
			// 	if angleDifference(currentBearing, edgeBearing) > trajectoryAngleTolerance {
			// 		continue
			// 	}
			// }

			var arrivalTime time.Time
			var departureTime time.Time
			var waitDuration time.Duration

			if edge.TripID == "transfer" {
				waitDuration = time.Duration(edge.Time) * time.Second
				if waitDuration > maxWaitBetweenTrips {
					continue
				}
				arrivalTime = currentArrival.Add(waitDuration)
				departureTime = arrivalTime
			} else {
				depSec := edge.DepartureTime
				day := currentArrival.Truncate(24 * time.Hour)
				departure := day.Add(time.Duration(depSec) * time.Second)
				if departure.Before(currentArrival) {
					departure = departure.Add(24 * time.Hour)
				}
				if departure.After(currentArrival) || departure.Equal(currentArrival) {
					arrivalTime = departure.Add(time.Duration(edge.Time) * time.Second)
					departureTime = departure
					waitDuration = departureTime.Sub(currentArrival)
					if waitDuration > maxWaitBetweenTrips {
						continue
					}
				} else {
					continue
				}
			}

			if arrivalTime.Sub(when) > maxTravelDuration {
				continue
			}

			currentTransfers := prev[current].Transfers
			lastTrip := prev[current].LastTrip
			newTransfers := currentTransfers
			var newLastTrip string
			if edge.TripID == "transfer" {
				newLastTrip = lastTrip
			} else {
				newLastTrip = edge.TripID
				if lastTrip != "" && lastTrip != edge.TripID {
					newTransfers++
				}
			}
			if newTransfers > maxTransfers {
				continue
			}

			costTime := arrivalTime
			if edge.TripID != "transfer" && lastTrip != "" && lastTrip != edge.TripID {
				costTime = costTime.Add(transferPenalty)
			}

			existingCost, hasCost := cost[edge.To]
			existingArrival := arrival[edge.To]
			existingInfo, havePrev := prev[edge.To]
			shouldRelax := !hasCost || existingCost.IsZero() || costTime.Before(existingCost)
			if !shouldRelax && costTime.Equal(existingCost) {
				if existingArrival.IsZero() || arrivalTime.Before(existingArrival) {
					shouldRelax = true
				} else if havePrev && arrivalTime.Equal(existingArrival) {
					if newTransfers < existingInfo.Transfers {
						shouldRelax = true
					} else if waitDuration < existingInfo.WaitDuration {
						shouldRelax = true
					}
				}
			}
			if shouldRelax {
				ancestor := current
				createsCycle := false
				for ancestor != "" {
					if ancestor == edge.To {
						createsCycle = true
						break
					}
					info, ok := prev[ancestor]
					if !ok {
						break
					}
					ancestor = info.Prev
				}
				if createsCycle {
					continue
				}
				arrival[edge.To] = arrivalTime
				cost[edge.To] = costTime
				prev[edge.To] = PathInfo{Prev: current, TripID: edge.TripID, DepartureTime: departureTime, Transfers: newTransfers, LastTrip: newLastTrip, WaitDuration: waitDuration}
				heap.Push(pq, &pqItem{Stop: edge.To, Cost: costTime})
			}
		}
	}
	return nil
}

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
}

type Edge struct {
	To            string
	TripID        string
	Time          float64
	DepartureTime float64
}

type TripDetail struct {
	RouteShortName string   `json:"route_short_name"`
	AgencyName     string   `json:"agency_name"`
	Stops          []string `json:"stops"`
}

type Leg struct {
	From          string        `json:"-"`
	FromStop      *types.Stop   `json:"from"`
	To            string        `json:"-"`
	ToStop        *types.Stop   `json:"to"`
	TripID        string        `json:"-"`
	Trip          *types.Trip   `json:"trip"`
	Stops         []string      `json:"stops"`
	Agency        *types.Agency `json:"agency"`
	Route         *types.Route  `json:"route"`
	DepartureTime time.Time     `json:"departure_time"`
	ArrivalTime   time.Time     `json:"arrival_time"`
}

type Route struct {
	Legs []Leg `json:"legs"`
}

func (r *Route) EndTime() time.Time {
	if len(r.Legs) == 0 {
		return time.Time{}
	}
	return r.Legs[len(r.Legs)-1].ArrivalTime
}

func (r *Route) StartTime() time.Time {
	if len(r.Legs) == 0 {
		return time.Time{}
	}
	return r.Legs[0].DepartureTime
}

func (r *Route) Duration() time.Duration {
	if len(r.Legs) == 0 {
		return 0
	}
	return r.Legs[len(r.Legs)-1].ArrivalTime.Sub(r.Legs[0].DepartureTime)
}

type PathInfo struct {
	Prev          string
	TripID        string
	DepartureTime time.Time
	Transfers     int
	LastTrip      string
	WaitDuration  time.Duration
}

func main() {
	tp := NewTripPlanner()
	wg := &sync.WaitGroup{}
	if err := tp.LoadData("data"); err != nil {
		fmt.Printf("Failed to load data: %v\n", err)
		os.Exit(1)
	}
	wg.Wait()
	if err := tp.Preprocess(); err != nil {
		fmt.Printf("Failed to preprocess data: %v\n", err)
		os.Exit(1)
	}

	http.HandleFunc("/api/stops", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "GET")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type")
		stopList := []types.Stop{}
		for _, s := range tp.stops {
			if _, hasConnections := tp.graph[s.StopID]; hasConnections {
				stopList = append(stopList, *s)
			}
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(stopList)
	})

	http.HandleFunc("/api/route", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "GET")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type")
		from := r.URL.Query().Get("from")
		to := r.URL.Query().Get("to")
		whenStr := r.URL.Query().Get("when")
		numStr := r.URL.Query().Get("num")
		num := 3
		if numStr != "" {
			if parsed, err := strconv.Atoi(numStr); err == nil && parsed > 0 {
				num = parsed
			}
		}
		when := time.Now()
		if whenStr != "" {
			if parsed, err := time.Parse(time.RFC3339, whenStr); err == nil {
				when = parsed
			}
		}
		if from == "" || to == "" {
			w.WriteHeader(http.StatusBadRequest)
			json.NewEncoder(w).Encode(map[string]string{"error": "from and to parameters required"})
			return
		}
		routes := tp.FindRoutes(from, to, when, num)
		if len(routes) == 0 {
			w.WriteHeader(http.StatusNotFound)
			json.NewEncoder(w).Encode(map[string]string{"error": "no route found"})
			return
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(routes)
	})
	log.Printf("Listening on 8080")
	http.ListenAndServe(":8080", nil)
}