trip-planner/internal/routing/graph.go

package routing

import "sort"

// Edge represents a graph edge connecting two nodes.
type Edge struct {
	From          *Node
	To            *Node
	Kind          EdgeKind
	Duration      int    // travel time in seconds
	Transport     string // transport type (train, plane, bus)
	TransportType string // deprecated: use Transport instead
	IsTransfer    bool   // whether this edge involves a transfer
	Departure     string // ISO 8601 departure time
	Arrival       string // ISO 8601 arrival time
}

// NodeType represents the type of a graph node.
type NodeType int

const (
	// NodeTypeStation represents a train station.
	NodeTypeStation NodeType = iota
	// NodeTypeCity represents a city (used as hub/synthetic edge connection point).
	NodeTypeCity
)

// Node represents a graph node (station or city).
type Node struct {
	ID       string
	Type     NodeType
	Name     string // display name (station title or city name)
	CityCode string // for stations, the city code they belong to
}

// EdgeKind represents the kind of edge in the graph.
type EdgeKind int

const (
	// EdgeKindReal represents a real scheduled trip (actual route segment).
	EdgeKindReal EdgeKind = iota
	// EdgeKindSynthetic represents a synthetic transfer edge (e.g., city↔airport).
	EdgeKindSynthetic
)

// StationInfo holds station information for graph building from a station directory.
type StationInfo struct {
	ID       string
	Name     string
	CityCode string
	CityName string
}

// Graph represents a routing graph with nodes (stations/cities) and edges (scheduled trips/transfers).
type Graph struct {
	nodes []*Node
	edges []*Edge
}

// NewGraph creates a new empty routing graph.
func NewGraph() *Graph {
	return &Graph{
		nodes: []*Node{},
		edges: []*Edge{},
	}
}

// AddNode adds a node to the graph.
func (g *Graph) AddNode(node *Node) {
	g.nodes = append(g.nodes, node)
}

// AddEdge adds an edge to the graph.
func (g *Graph) AddEdge(edge *Edge) {
	g.edges = append(g.edges, edge)
}

// Nodes returns all nodes in the graph.
func (g *Graph) Nodes() []*Node {
	result := make([]*Node, len(g.nodes))
	copy(result, g.nodes)
	return result
}

// Edges returns all edges in the graph.
func (g *Graph) Edges() []*Edge {
	result := make([]*Edge, len(g.edges))
	copy(result, g.edges)
	return result
}

// BuildGraphFromStations builds a routing graph from a list of station info records.
// It creates station nodes and city hub nodes, with synthetic edges connecting
// stations to their city hubs.
func BuildGraphFromStations(stations []StationInfo) *Graph {
	graph := NewGraph()

	// Track city nodes by code to avoid duplicates
	cityNodes := make(map[string]*Node)

	// Add all station nodes and create/connect city hub nodes
	for _, si := range stations {
		// Add station node
		station := &Node{
			ID:       si.ID,
			Type:     NodeTypeStation,
			Name:     si.Name,
			CityCode: si.CityCode,
		}
		graph.AddNode(station)

		// Create or retrieve city hub node
		cityKey := "city:" + si.CityCode
		if _, exists := cityNodes[si.CityCode]; !exists {
			cityNode := &Node{
				ID:   cityKey,
				Type: NodeTypeCity,
				Name: si.CityName,
			}
			graph.AddNode(cityNode)
			cityNodes[si.CityCode] = cityNode
		}

		// Add synthetic edge: station <-> city hub
		cityNode := cityNodes[si.CityCode]
		graph.AddEdge(&Edge{
			From:       station,
			To:         cityNode,
			Kind:       EdgeKindSynthetic,
			Duration:   300, // 5 min synthetic transfer
			Transport:  "train",
			IsTransfer: true,
		})

		// Add reverse synthetic edge: city hub -> station
		graph.AddEdge(&Edge{
			From:       cityNode,
			To:         station,
			Kind:       EdgeKindSynthetic,
			Duration:   300, // 5 min synthetic transfer
			Transport:  "train",
			IsTransfer: true,
		})
	}

	return graph
}

// SortEdges sorts edges by duration in ascending order (shortest first).
func SortEdges(edges []*Edge) {
	sort.Slice(edges, func(i, j int) bool {
		return edges[i].Duration < edges[j].Duration
	})
}

// buildAdjacencyList builds an adjacency list from the graph's edges.
func (g *Graph) buildAdjacencyList() map[string][]*Edge {
	adj := make(map[string][]*Edge)
	for _, edge := range g.edges {
		adj[edge.From.ID] = append(adj[edge.From.ID], edge)
	}
	return adj
}

// nodesByID returns a node by its ID from the graph's nodes.
func (g *Graph) nodesByID(id string) *Node {
	for _, n := range g.nodes {
		if n.ID == id {
			return n
		}
	}
	return nil
}

// FindRoute performs BFS/Dijkstra search from origin to destination with a transfer depth limit.
// It returns the best itinerary found within the transfer limit.
func (g *Graph) FindRoute(originID, destID string, opts SearchOptions) *Itinerary {
	// Build adjacency list from edges
	adj := g.buildAdjacencyList()

	// BFS with transfer tracking
	// State: (nodeID, transfersUsed, accumulatedDuration, lastArrivalTime, path)
	startNode := g.nodesByID(originID)
	destNode := g.nodesByID(destID)

	if startNode == nil || destNode == nil {
		return nil
	}

	// Queue for BFS: each element is a state
	type bfsState struct {
		nodeID      string
		transfers   int
		duration    int
		lastArrival string // arrival time at current node (for MCT calculation)
		itinerary   *Itinerary
	}

	// Track the minimum transfers seen for each node to prune suboptimal paths
	visited := make(map[string]int) // nodeID -> min transfers seen

	// Initialize with the start node
	initial := bfsState{
		nodeID:      originID,
		transfers:   0,
		duration:    0,
		lastArrival: "",
		itinerary:   &Itinerary{Legs: []RouteLeg{}},
	}

	// Use a simple slice as priority queue - sort by (duration, transfers)
	var queue []bfsState
	queue = append(queue, initial)

	var best *Itinerary

	for len(queue) > 0 {
		// Pop the state with shortest duration (and fewest transfers as tiebreaker)
		current := queue[0]
		queue = queue[1:]

		// If we've reached the destination, potentially update best result
		if current.nodeID == destID {
			if best == nil || current.duration < best.TotalDuration ||
				(current.duration == best.TotalDuration && current.transfers < best.TotalTransfers) {
				best = current.itinerary
				// Recalculate best metrics from legs
				best.TotalDuration = current.duration
				best.TotalTransfers = current.transfers
			}
			// Don't continue from destination - we've arrived
			continue
		}

		// Prune if we've exceeded max transfers
		if opts.MaxTransfers >= 0 && current.transfers >= opts.MaxTransfers {
			continue
		}

		// Explore outgoing edges
		for _, edge := range adj[current.nodeID] {
			nextNode := edge.To

			// Calculate new duration
			newDuration := current.duration + edge.Duration

			// Calculate transfer time if this is not the first leg
			transferTime := 0
			if current.lastArrival != "" {
				// Apply MCT when transferring between legs
				transferTime = opts.MCT
			}

			newDurationWithMCT := newDuration + transferTime

			// Check if we've visited this node with fewer or equal transfers
			visKey := current.nodeID
			if existingTransfers, ok := visited[visKey]; ok {
				if current.transfers+1 >= existingTransfers {
					// Already visited this node with fewer or equal transfers, skip
					continue
				}
			}
			visited[visKey] = current.transfers + 1

			newTransfers := current.transfers
			if edge.IsTransfer {
				newTransfers++
			}

			// Build new itinerary legs
			newLegs := make([]RouteLeg, len(current.itinerary.Legs)+1)
			copy(newLegs, current.itinerary.Legs)

			// First leg: From is the origin node, subsequent legs use the previous edge's To
			if len(current.itinerary.Legs) == 0 {
				newLegs[len(current.itinerary.Legs)] = RouteLeg{
					From:       g.nodesByID(originID), // origin node as From
					To:         nextNode,
					Duration:   edge.Duration,
					Transport:  edge.Transport,
					IsTransfer: edge.IsTransfer,
				}
			} else {
				newLegs[len(current.itinerary.Legs)] = RouteLeg{
					From:       current.itinerary.Legs[len(current.itinerary.Legs)-1].To,
					To:         nextNode,
					Duration:   edge.Duration,
					Transport:  edge.Transport,
					IsTransfer: edge.IsTransfer,
				}
			}

			newItinerary := &Itinerary{
				Legs:           newLegs,
				TotalDuration:  newDurationWithMCT,
				TotalTransfers: newTransfers,
			}

			queue = append(queue, bfsState{
				nodeID:      nextNode.ID,
				transfers:   newTransfers,
				duration:    newDurationWithMCT,
				lastArrival: edge.Arrival, // arrival time at next node
				itinerary:   newItinerary,
			})
		}

		// Re-sort queue by (duration, transfers) for priority
		sort.Slice(queue, func(i, j int) bool {
			if queue[i].duration != queue[j].duration {
				return queue[i].duration < queue[j].duration
			}
			return queue[i].transfers < queue[j].transfers
		})
	}

	if best == nil {
		return nil
	}
	return best
}

// ApplyMCT applies Minimum Connection Time rules to the itinerary.
// It adjusts transfer times based on node types, city tiers, and check-in requirements.
func (g *Graph) ApplyMCT(itinerary *Itinerary, mctBase int) *Itinerary {
	if itinerary == nil || len(itinerary.Legs) <= 1 {
		// No transfers needed, return as-is
		return itinerary
	}

	// MCT base default: 30 minutes (1800 seconds)
	if mctBase <= 0 {
		mctBase = 1800
	}

	// Create a working copy of legs
	adjustedLegs := make([]RouteLeg, len(itinerary.Legs))
	copy(adjustedLegs, itinerary.Legs)

	for i := 1; i < len(adjustedLegs); i++ {
		prevLeg := &adjustedLegs[i-1]
		currLeg := &adjustedLegs[i]

		// Determine MCT based on node types and transfer kinds
		mct := mctBase

		// Reduce MCT for city hub transfers (the transfer point node is a city)
		// The transfer point is the destination of the previous leg / start of current leg
		transferPoint := prevLeg.To // = currLeg.From
		if transferPoint.Type == NodeTypeCity {
			mct = mctBase / 2 // 30 min -> 15 min for city hub transfers
		}

		// Increase MCT for mode changes (different transport types)
		if prevLeg.Transport != currLeg.Transport {
			mct = mctBase + 600 // 30 min + 10 min for mode change
		}

		// Add the MCT to the total duration (as waiting time at transfer)
		itinerary.TotalDuration += mct
	}

	// Recalculate leg structure with proper transfer timing
	itinerary.Legs = adjustedLegs
	return itinerary
}

// SearchOptions configures the route search behavior.
type SearchOptions struct {
	// MaxTransfers limits the number of transfers allowed in the route.
	MaxTransfers int
	// MCT is the minimum connection time in seconds at transfer points.
	MCT int
	// FarTerm indicates if the search date is far-term (affects caching/TTL).
	FarTerm bool
}

// Itinerary represents a complete route with legs and summary metrics.
type Itinerary struct {
	Legs           []RouteLeg
	TotalDuration  int // total travel time in seconds
	TotalTransfers int // number of transfers
	Cost           int // cost in minor currency units (e.g., rubles)
	// Identifier for the route (e.g., search_id + route_id)
	ID string
}

// RouteLeg represents a single leg of a route (one edge between two nodes).
type RouteLeg struct {
	From       *Node
	To         *Node
	Departure  string // ISO 8601 departure time
	Arrival    string // ISO 8601 arrival time
	Duration   int    // travel time in seconds
	Transport  string // transport type (train, plane, bus)
	IsTransfer bool
}

// SearchResult represents the result of a route search.
type SearchResult struct {
	// Itineraries are the found routes, sorted by Pareto ranking (time, transfers, cost).
	Itineraries []*Itinerary
	// SearchMetadata contains information about the search execution.
	Metadata map[string]interface{}
}

// FindRoutesPareto finds Pareto-optimal routes (time, transfers, cost) from origin to destination.
// It runs the search algorithm and returns multiple routes that are not dominated by any other
// route in all three metrics simultaneously.
func (g *Graph) FindRoutesPareto(originID, destID string, opts SearchOptions) []*Itinerary {
	// Run multiple searches with different strategies to find diverse routes
	var allItineraries []*Itinerary

	// Search with different max transfer limits to find diverse routes
	for maxTransfers := 0; maxTransfers <= opts.MaxTransfers; maxTransfers++ {
		optsCopy := opts
		optsCopy.MaxTransfers = maxTransfers

		result := g.FindRoute(originID, destID, optsCopy)
		if result != nil && result.TotalDuration > 0 {
			allItineraries = append(allItineraries, result)
		}
	}

	// Sort by total duration (primary), then transfers (secondary), then cost (tertiary)
	sort.Slice(allItineraries, func(i, j int) bool {
		if allItineraries[i].TotalDuration != allItineraries[j].TotalDuration {
			return allItineraries[i].TotalDuration < allItineraries[j].TotalDuration
		}
		if allItineraries[i].TotalTransfers != allItineraries[j].TotalTransfers {
			return allItineraries[i].TotalTransfers < allItineraries[j].TotalTransfers
		}
		return allItineraries[i].Cost < allItineraries[j].Cost
	})

	// Pareto filter: remove dominated routes
	// A route is dominated if another route is better or equal in all metrics (time, transfers, cost)
	var pareto []*Itinerary
	for _, candidate := range allItineraries {
		dominated := false
		for _, existing := range pareto {
			// Check if existing dominates candidate
			if existing.TotalDuration <= candidate.TotalDuration &&
				existing.TotalTransfers <= candidate.TotalTransfers &&
				existing.Cost <= candidate.Cost &&
				(existing.TotalDuration < candidate.TotalDuration ||
					existing.TotalTransfers < candidate.TotalTransfers ||
					existing.Cost < candidate.Cost) {
				dominated = true
				break
			}
		}
		if !dominated {
			pareto = append(pareto, candidate)
		}
	}

	return pareto
}