Sportstime/SportsTime/Planning/Engine/GameDAGRouter.swift

//
//  GameDAGRouter.swift
//  SportsTime
//
//  Time-expanded DAG + Beam Search algorithm for route finding.
//
//  Key insight: This is NOT "which subset of N games should I attend?"
//  This IS: "what time-respecting paths exist through a graph of games?"
//
//  The algorithm:
//    1. Bucket games by calendar day
//    2. Build directed edges where time moves forward AND driving is feasible
//    3. Beam search: keep top K paths at each depth
//    4. Dominance pruning: discard inferior paths
//
//  Complexity: O(days × beamWidth × avgNeighbors) ≈ 900 operations for 5-day, 78-game scenario
//  (vs 2^78 for naive subset enumeration)
//

import Foundation
import CoreLocation

enum GameDAGRouter {

    // MARK: - Configuration

    /// Default beam width - how many partial routes to keep at each step
    /// Increased to ensure we preserve diverse route lengths (short and long trips)
    private static let defaultBeamWidth = 50

    /// Maximum options to return (increased to provide more diverse trip lengths)
    private static let maxOptions = 50

    /// Buffer time after game ends before we can depart (hours)
    private static let gameEndBufferHours: Double = 3.0

    /// Maximum days ahead to consider for next game (1 = next day only, 5 = allows multi-day drives)
    private static let maxDayLookahead = 5

    // MARK: - Public API

    /// Finds best routes through the game graph using DAG + beam search.
    ///
    /// This replaces the exponential GeographicRouteExplorer with a polynomial-time algorithm.
    ///
    /// - Parameters:
    ///   - games: All games to consider, in any order (will be sorted internally)
    ///   - stadiums: Dictionary mapping stadium IDs to Stadium objects
    ///   - constraints: Driving constraints (number of drivers, max hours per day)
    ///   - anchorGameIds: Games that MUST appear in every valid route (for Scenario B)
    ///   - allowRepeatCities: If false, each city can only appear once in a route
    ///   - beamWidth: How many partial routes to keep at each depth (default 30)
    ///
    /// - Returns: Array of valid game combinations, sorted by score (most games, least driving)
    ///
    static func findRoutes(
        games: [Game],
        stadiums: [UUID: Stadium],
        constraints: DrivingConstraints,
        anchorGameIds: Set<UUID> = [],
        allowRepeatCities: Bool = true,
        beamWidth: Int = defaultBeamWidth
    ) -> [[Game]] {

        // Edge cases
        guard !games.isEmpty else { return [] }
        if games.count == 1 {
            // Single game - just return it if it satisfies anchors
            if anchorGameIds.isEmpty || anchorGameIds.contains(games[0].id) {
                return [games]
            }
            return []
        }
        if games.count == 2 {
            // Two games - check if both are reachable
            let sorted = games.sorted { $0.startTime < $1.startTime }
            if canTransition(from: sorted[0], to: sorted[1], stadiums: stadiums, constraints: constraints) {
                if anchorGameIds.isSubset(of: Set(sorted.map { $0.id })) {
                    return [sorted]
                }
            }
            // Can't connect them - return individual games if they satisfy anchors
            if anchorGameIds.isEmpty {
                return [[sorted[0]], [sorted[1]]]
            }
            return []
        }

        // Step 1: Sort games chronologically
        let sortedGames = games.sorted { $0.startTime < $1.startTime }

        // Step 2: Bucket games by calendar day
        let buckets = bucketByDay(games: sortedGames)
        let sortedDays = buckets.keys.sorted()

        guard !sortedDays.isEmpty else { return [] }


        // Step 3: Initialize beam with first day's games
        var beam: [[Game]] = []
        if let firstDayGames = buckets[sortedDays[0]] {
            for game in firstDayGames {
                beam.append([game])
            }
        }

        // Also include option to skip first day entirely and start later
        // (handled by having multiple starting points in beam)
        for dayIndex in sortedDays.dropFirst().prefix(maxDayLookahead - 1) {
            if let dayGames = buckets[dayIndex] {
                for game in dayGames {
                    beam.append([game])
                }
            }
        }


        // Step 4: Expand beam day by day
        for (_, dayIndex) in sortedDays.dropFirst().enumerated() {
            let todaysGames = buckets[dayIndex] ?? []
            var nextBeam: [[Game]] = []

            for path in beam {
                guard let lastGame = path.last else { continue }
                let lastGameDay = dayIndexFor(lastGame.startTime, referenceDate: sortedGames[0].startTime)

                // Only consider games on this day or within lookahead
                if dayIndex > lastGameDay + maxDayLookahead {
                    // This path is too far behind, keep it as-is
                    nextBeam.append(path)
                    continue
                }

                // Try adding each of today's games
                for candidate in todaysGames {
                    // Check for repeat city violation during route building
                    if !allowRepeatCities {
                        let candidateCity = stadiums[candidate.stadiumId]?.city ?? ""
                        let pathCities = Set(path.compactMap { stadiums[$0.stadiumId]?.city })
                        if pathCities.contains(candidateCity) {
                            continue // Skip - would violate allowRepeatCities
                        }
                    }

                    if canTransition(from: lastGame, to: candidate, stadiums: stadiums, constraints: constraints) {
                        let newPath = path + [candidate]
                        nextBeam.append(newPath)
                    }
                }

                // Also keep the path without adding a game today (allows off-days)
                nextBeam.append(path)
            }

            // Dominance pruning + beam truncation
            beam = pruneAndTruncate(nextBeam, beamWidth: beamWidth, stadiums: stadiums)
        }

        // Step 5: Filter routes that contain all anchors
        let routesWithAnchors = beam.filter { path in
            let pathGameIds = Set(path.map { $0.id })
            return anchorGameIds.isSubset(of: pathGameIds)
        }

        // Step 6: Ensure geographic diversity in results
        // Group routes by their primary region (city with most games)
        // Then pick the best route from each region
        let finalRoutes = selectDiverseRoutes(routesWithAnchors, stadiums: stadiums, maxCount: maxOptions)

        print("🔍 DAG: Input games=\(games.count), beam final=\(beam.count), withAnchors=\(routesWithAnchors.count), final=\(finalRoutes.count)")
        if let best = finalRoutes.first {
            print("🔍 DAG: Best route has \(best.count) games")
        }

        return finalRoutes
    }

    /// Compatibility wrapper that matches GeographicRouteExplorer's interface.
    /// This allows drop-in replacement in ScenarioAPlanner and ScenarioBPlanner.
    static func findAllSensibleRoutes(
        from games: [Game],
        stadiums: [UUID: Stadium],
        anchorGameIds: Set<UUID> = [],
        allowRepeatCities: Bool = true,
        stopBuilder: ([Game], [UUID: Stadium]) -> [ItineraryStop]
    ) -> [[Game]] {
        // Use default driving constraints
        let constraints = DrivingConstraints.default

        return findRoutes(
            games: games,
            stadiums: stadiums,
            constraints: constraints,
            anchorGameIds: anchorGameIds,
            allowRepeatCities: allowRepeatCities
        )
    }

    // MARK: - Day Bucketing

    /// Groups games by calendar day index (0 = first day of trip, 1 = second day, etc.)
    private static func bucketByDay(games: [Game]) -> [Int: [Game]] {
        guard let firstGame = games.first else { return [:] }
        let referenceDate = firstGame.startTime

        var buckets: [Int: [Game]] = [:]
        for game in games {
            let dayIndex = dayIndexFor(game.startTime, referenceDate: referenceDate)
            buckets[dayIndex, default: []].append(game)
        }
        return buckets
    }

    /// Calculates the day index for a date relative to a reference date.
    private static func dayIndexFor(_ date: Date, referenceDate: Date) -> Int {
        let calendar = Calendar.current
        let refDay = calendar.startOfDay(for: referenceDate)
        let dateDay = calendar.startOfDay(for: date)
        let components = calendar.dateComponents([.day], from: refDay, to: dateDay)
        return components.day ?? 0
    }

    // MARK: - Transition Feasibility

    /// Determines if we can travel from game A to game B.
    ///
    /// Requirements:
    ///   1. B starts after A (time moves forward)
    ///   2. We have enough days between games to complete the drive
    ///   3. We can arrive at B before B starts
    ///
    private static func canTransition(
        from: Game,
        to: Game,
        stadiums: [UUID: Stadium],
        constraints: DrivingConstraints
    ) -> Bool {
        // Time must move forward
        guard to.startTime > from.startTime else {
            return false
        }

        // Same stadium = always feasible (no driving needed)
        if from.stadiumId == to.stadiumId { return true }

        // Get stadiums
        guard let fromStadium = stadiums[from.stadiumId],
              let toStadium = stadiums[to.stadiumId] else {
            // Missing stadium info - can't calculate distance, reject to be safe
            print("⚠️ DAG: Stadium lookup failed - from:\(stadiums[from.stadiumId] != nil) to:\(stadiums[to.stadiumId] != nil)")
            return false
        }

        let fromCoord = fromStadium.coordinate
        let toCoord = toStadium.coordinate

        // Calculate driving time
        let distanceMiles = TravelEstimator.haversineDistanceMiles(
            from: CLLocationCoordinate2D(latitude: fromCoord.latitude, longitude: fromCoord.longitude),
            to: CLLocationCoordinate2D(latitude: toCoord.latitude, longitude: toCoord.longitude)
        ) * 1.3  // Road routing factor

        let drivingHours = distanceMiles / 60.0  // Average 60 mph

        // Calculate available driving time between games
        // After game A ends (+ buffer), how much time until game B starts (- buffer)?
        let departureTime = from.startTime.addingTimeInterval(gameEndBufferHours * 3600)
        let deadline = to.startTime.addingTimeInterval(-3600)  // 1 hour buffer before game
        let availableSeconds = deadline.timeIntervalSince(departureTime)
        let availableHours = availableSeconds / 3600.0

        // Calculate how many driving days we have
        // Each day can have maxDailyDrivingHours of driving
        let calendar = Calendar.current
        let fromDay = calendar.startOfDay(for: from.startTime)
        let toDay = calendar.startOfDay(for: to.startTime)
        let daysBetween = calendar.dateComponents([.day], from: fromDay, to: toDay).day ?? 0

        // Available driving hours = days between * max per day
        // (If games are same day, daysBetween = 0, but we might still have hours available)
        let maxDrivingHoursAvailable: Double
        if daysBetween == 0 {
            // Same day - only have hours between games
            maxDrivingHoursAvailable = max(0, availableHours)
        } else {
            // Multi-day - can drive each day
            maxDrivingHoursAvailable = Double(daysBetween) * constraints.maxDailyDrivingHours
        }

        // Check if we have enough driving time
        guard drivingHours <= maxDrivingHoursAvailable else {
            return false
        }

        // Also verify we can arrive before game starts (sanity check)
        guard availableHours >= drivingHours else {
            return false
        }

        return true
    }

    // MARK: - Geographic Diversity

    /// Selects diverse routes from the candidate set.
    /// Ensures diversity by BOTH route length (city count) AND primary city.
    /// This guarantees users see 2-city trips alongside 5+ city trips.
    private static func selectDiverseRoutes(
        _ routes: [[Game]],
        stadiums: [UUID: Stadium],
        maxCount: Int
    ) -> [[Game]] {
        guard !routes.isEmpty else { return [] }

        // Group routes by city count (route length)
        var routesByLength: [Int: [[Game]]] = [:]
        for route in routes {
            let cityCount = Set(route.compactMap { stadiums[$0.stadiumId]?.city }).count
            routesByLength[cityCount, default: []].append(route)
        }

        // Sort routes within each length by score
        for (length, lengthRoutes) in routesByLength {
            routesByLength[length] = lengthRoutes.sorted {
                scorePath($0, stadiums: stadiums) > scorePath($1, stadiums: stadiums)
            }
        }

        // Allocate slots to each length category
        // Goal: ensure at least 1 route per length category if available
        let sortedLengths = routesByLength.keys.sorted()
        let minPerLength = max(1, maxCount / max(1, sortedLengths.count))

        var selectedRoutes: [[Game]] = []
        var selectedIds = Set<String>()

        // First pass: take best route(s) from each length category
        for length in sortedLengths {
            if selectedRoutes.count >= maxCount { break }
            if let lengthRoutes = routesByLength[length] {
                let toTake = min(minPerLength, lengthRoutes.count, maxCount - selectedRoutes.count)
                for route in lengthRoutes.prefix(toTake) {
                    let key = route.map { $0.id.uuidString }.joined(separator: "-")
                    if !selectedIds.contains(key) {
                        selectedRoutes.append(route)
                        selectedIds.insert(key)
                    }
                }
            }
        }

        // Second pass: fill remaining slots, prioritizing geographic diversity
        if selectedRoutes.count < maxCount {
            // Group remaining routes by primary city
            var remainingByCity: [String: [[Game]]] = [:]
            for route in routes {
                let key = route.map { $0.id.uuidString }.joined(separator: "-")
                if !selectedIds.contains(key) {
                    let city = getPrimaryCity(for: route, stadiums: stadiums)
                    remainingByCity[city, default: []].append(route)
                }
            }

            // Sort by score within each city
            for (city, cityRoutes) in remainingByCity {
                remainingByCity[city] = cityRoutes.sorted {
                    scorePath($0, stadiums: stadiums) > scorePath($1, stadiums: stadiums)
                }
            }

            // Round-robin from each city
            let sortedCities = remainingByCity.keys.sorted { city1, city2 in
                let score1 = remainingByCity[city1]?.first.map { scorePath($0, stadiums: stadiums) } ?? 0
                let score2 = remainingByCity[city2]?.first.map { scorePath($0, stadiums: stadiums) } ?? 0
                return score1 > score2
            }

            var cityIndices: [String: Int] = [:]
            while selectedRoutes.count < maxCount {
                var addedAny = false
                for city in sortedCities {
                    if selectedRoutes.count >= maxCount { break }
                    let idx = cityIndices[city] ?? 0
                    if let cityRoutes = remainingByCity[city], idx < cityRoutes.count {
                        let route = cityRoutes[idx]
                        let key = route.map { $0.id.uuidString }.joined(separator: "-")
                        if !selectedIds.contains(key) {
                            selectedRoutes.append(route)
                            selectedIds.insert(key)
                            addedAny = true
                        }
                        cityIndices[city] = idx + 1
                    }
                }
                if !addedAny { break }
            }
        }

        return selectedRoutes
    }

    /// Gets the primary city for a route (where most games are played).
    private static func getPrimaryCity(for route: [Game], stadiums: [UUID: Stadium]) -> String {
        var cityCounts: [String: Int] = [:]
        for game in route {
            let city = stadiums[game.stadiumId]?.city ?? "Unknown"
            cityCounts[city, default: 0] += 1
        }
        return cityCounts.max(by: { $0.value < $1.value })?.key ?? "Unknown"
    }

    // MARK: - Scoring and Pruning

    /// Scores a path. Higher = better.
    /// Prefers: more games, less driving, geographic coherence
    private static func scorePath(_ path: [Game], stadiums: [UUID: Stadium]) -> Double {
        // Handle empty or single-game paths
        guard path.count > 1 else {
            return Double(path.count) * 100.0
        }

        let gameCount = Double(path.count)

        // Calculate total driving
        var totalDriving: Double = 0
        for i in 0..<(path.count - 1) {
            totalDriving += estimateDrivingHours(from: path[i], to: path[i + 1], stadiums: stadiums)
        }

        // Score: heavily weight game count, penalize driving
        return gameCount * 100.0 - totalDriving * 2.0
    }

    /// Estimates driving hours between two games.
    private static func estimateDrivingHours(
        from: Game,
        to: Game,
        stadiums: [UUID: Stadium]
    ) -> Double {
        // Same stadium = 0 driving
        if from.stadiumId == to.stadiumId { return 0 }

        guard let fromStadium = stadiums[from.stadiumId],
              let toStadium = stadiums[to.stadiumId] else {
            return 5.0  // Fallback: assume 5 hours
        }

        let fromCoord = fromStadium.coordinate
        let toCoord = toStadium.coordinate

        let distanceMiles = TravelEstimator.haversineDistanceMiles(
            from: CLLocationCoordinate2D(latitude: fromCoord.latitude, longitude: fromCoord.longitude),
            to: CLLocationCoordinate2D(latitude: toCoord.latitude, longitude: toCoord.longitude)
        ) * 1.3

        return distanceMiles / 60.0
    }

    /// Prunes dominated paths and truncates to beam width.
    /// Maintains diversity by both ending city AND route length to ensure short trips aren't eliminated.
    private static func pruneAndTruncate(
        _ paths: [[Game]],
        beamWidth: Int,
        stadiums: [UUID: Stadium]
    ) -> [[Game]] {
        // Remove exact duplicates
        var uniquePaths: [[Game]] = []
        var seen = Set<String>()

        for path in paths {
            let key = path.map { $0.id.uuidString }.joined(separator: "-")
            if !seen.contains(key) {
                seen.insert(key)
                uniquePaths.append(path)
            }
        }

        // Group paths by unique city count (route length)
        // This ensures we keep short trips (2 cities) alongside long trips (5+ cities)
        var pathsByLength: [Int: [[Game]]] = [:]
        for path in uniquePaths {
            let cityCount = Set(path.compactMap { stadiums[$0.stadiumId]?.city }).count
            pathsByLength[cityCount, default: []].append(path)
        }

        // Sort paths within each length group by score
        for (length, lengthPaths) in pathsByLength {
            pathsByLength[length] = lengthPaths.sorted {
                scorePath($0, stadiums: stadiums) > scorePath($1, stadiums: stadiums)
            }
        }

        // Allocate beam slots proportionally to length groups, with minimum per group
        let sortedLengths = pathsByLength.keys.sorted()
        let minPerLength = max(2, beamWidth / max(1, sortedLengths.count))

        var pruned: [[Game]] = []

        // First pass: take minimum from each length group
        for length in sortedLengths {
            if let lengthPaths = pathsByLength[length] {
                let toTake = min(minPerLength, lengthPaths.count)
                pruned.append(contentsOf: lengthPaths.prefix(toTake))
            }
        }

        // Second pass: fill remaining slots with best paths overall
        if pruned.count < beamWidth {
            let remaining = beamWidth - pruned.count
            let prunedIds = Set(pruned.map { $0.map { $0.id.uuidString }.joined(separator: "-") })

            // Get all paths not yet added, sorted by score
            var additional = uniquePaths.filter {
                !prunedIds.contains($0.map { $0.id.uuidString }.joined(separator: "-"))
            }
            additional.sort { scorePath($0, stadiums: stadiums) > scorePath($1, stadiums: stadiums) }

            pruned.append(contentsOf: additional.prefix(remaining))
        }

        // Final truncation
        return Array(pruned.prefix(beamWidth))
    }
}