go-gtfs/cmd/planner/main.go

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"
	"time"

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

type Edge struct {
	To            string
	TripID        string
	Time          types.SecondsAfterMidnight
	DepartureTime types.SecondsAfterMidnight
}

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

type JSONTrip struct {
	TripId        string `json:"trip_id"`
	RouteId       string `json:"route_id"`
	AgencyName    string `json:"agency_name"`
	TripHeadsign  string `json:"trip_headsign"`
	TripShortName string `json:"trip_short_name"`
}

type Leg struct {
	From *types.StopTime `json:"start"`
	Trip *JSONTrip       `json:"trip"`
	To   *types.StopTime `json:"end"`
}

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

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

func (r *Route) StartTime() types.SecondsAfterMidnight {
	if len(r.Legs) == 0 {
		return 0
	}
	return r.Legs[0].From.DepartureTime
}

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

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

func main() {
	tripData, err := reader.LoadTripData("data")
	if err != nil {
		log.Fatalf("unable to load data %v", err)
	}
	tp := NewTripPlanner(tripData)

	if err := tp.Preprocess(); err != nil {
		fmt.Printf("Failed to preprocess data: %v\n", err)
		os.Exit(1)
	}

	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.ArrivalTime))
			}
		}

	}

	http.HandleFunc("/api/stops", func(w http.ResponseWriter, r *http.Request) {

		stopList := []types.Stop{}
		for _, s := range tp.Stops {
			if len(s.Trips) > 0 {
				stopList = append(stopList, *s)
			}
		}
		w.WriteHeader(http.StatusOK)
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(stopList)
	})

	http.HandleFunc("/api/trips", func(w http.ResponseWriter, r *http.Request) {
		from := r.URL.Query().Get("from")
		whenStr := r.URL.Query().Get("when")
		if from == "" {
			w.WriteHeader(http.StatusBadRequest)
			json.NewEncoder(w).Encode(map[string]string{"error": "from parameter required"})
			return
		}
		stop, ok := tp.Stops[from]
		if !ok {
			w.WriteHeader(http.StatusNotFound)
			json.NewEncoder(w).Encode(map[string]string{"error": "stop not found"})
			return
		}
		when := time.Now()
		if whenStr != "" {
			var err error
			when, err = time.Parse(time.RFC3339, whenStr)
			if err != nil {
				w.WriteHeader(http.StatusBadRequest)
				json.NewEncoder(w).Encode(map[string]string{"error": "invalid when parameter"})
				return
			}
		}
		trips := []map[string]interface{}{}
		for trip := range stop.GetTripsAfter(when) {
			// Find the index of the stop in the trip
			startIdx := 0
			for i, st := range trip.Stops {
				if st.StopId == from {
					startIdx = i
					break
				}
			}
			agencyName := ""
			if trip.Agency != nil {
				agencyName = trip.Agency.AgencyName
			}
			tripData := map[string]interface{}{
				"trip_id":     trip.TripId,
				"headsign":    trip.TripHeadsign,
				"short_name":  trip.TripShortName,
				"route_id":    trip.RouteId,
				"agency_id":   trip.AgencyId,
				"agency_name": agencyName,
				"stops":       []map[string]interface{}{},
			}
			for _, st := range trip.Stops[startIdx:] {
				tripData["stops"] = append(tripData["stops"].([]map[string]interface{}), map[string]interface{}{
					"stop_id":        st.StopId,
					"stop_name":      st.Stop.StopName,
					"location":       []float64{st.Stop.StopLat, st.Stop.StopLon},
					"arrival_time":   st.ArrivalTime,
					"departure_time": st.DepartureTime,
				})
			}
			trips = append(trips, tripData)
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(trips)
	})

	http.HandleFunc("/api/route", func(w http.ResponseWriter, r *http.Request) {

		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.DateTime, 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
		}
		log.Printf("using num %v", num)

		log.Printf("start time %v", when)
		route, err := tp.FindRoute(from, to, when)
		if err != nil {
			w.WriteHeader(http.StatusNotFound)
			json.NewEncoder(w).Encode(map[string]string{"error": "no route found"})
			return
		}
		w.WriteHeader(http.StatusOK)
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(route)
	})
	log.Printf("Listening on 8080")
	http.ListenAndServe(":8080", nil)
}