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
    private static let defaultBeamWidth = 30

    /// Maximum options to return
    private static let maxOptions = 10

    /// 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, 2 = allows one off-day)
    private static let maxDayLookahead = 2

    // 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)
    ///   - 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> = [],
        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 [] }

        print("[GameDAGRouter] \(games.count) games across \(sortedDays.count) days")
        print("[GameDAGRouter] Games per day: \(sortedDays.map { buckets[$0]?.count ?? 0 })")

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

        print("[GameDAGRouter] Initial beam size: \(beam.count)")

        // Step 4: Expand beam day by day
        for (index, 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
                }

                var addedAny = false

                // Try adding each of today's games
                for candidate in todaysGames {
                    if canTransition(from: lastGame, to: candidate, stadiums: stadiums, constraints: constraints) {
                        let newPath = path + [candidate]
                        nextBeam.append(newPath)
                        addedAny = true
                    }
                }

                // 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)
            print("[GameDAGRouter] Day \(dayIndex): nextBeam=\(nextBeam.count), after prune=\(beam.count), max games=\(beam.map { $0.count }.max() ?? 0)")
        }

        // 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 diverseRoutes = selectDiverseRoutes(routesWithAnchors, stadiums: stadiums, maxCount: maxOptions)

        print("[GameDAGRouter] Found \(routesWithAnchors.count) routes with anchors, returning \(diverseRoutes.count) diverse routes")
        for (i, route) in diverseRoutes.prefix(5).enumerated() {
            let cities = route.compactMap { stadiums[$0.stadiumId]?.city }.joined(separator: " → ")
            print("[GameDAGRouter] Route \(i+1): \(route.count) games - \(cities)")
        }

        return diverseRoutes
    }

    /// 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> = [],
        stopBuilder: ([Game], [UUID: Stadium]) -> [ItineraryStop]
    ) -> [[Game]] {
        // Use default driving constraints
        let constraints = DrivingConstraints.default

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

    // 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. Driving time is within daily limit
    ///   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 - use generous fallback
            // Assume 300 miles at 60 mph = 5 hours, which is usually feasible
            return true
        }

        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

        // Must be within daily limit
        guard drivingHours <= constraints.maxDailyDrivingHours else { return false }

        // Calculate if we can arrive in time
        let departureTime = from.startTime.addingTimeInterval(gameEndBufferHours * 3600)
        let arrivalTime = departureTime.addingTimeInterval(drivingHours * 3600)

        // Must arrive before game starts (with 1 hour buffer)
        let deadline = to.startTime.addingTimeInterval(-3600)
        guard arrivalTime <= deadline else { return false }

        return true
    }

    // MARK: - Geographic Diversity

    /// Selects geographically diverse routes from the candidate set.
    /// Groups routes by their primary city (where most games are) and picks the best from each region.
    private static func selectDiverseRoutes(
        _ routes: [[Game]],
        stadiums: [UUID: Stadium],
        maxCount: Int
    ) -> [[Game]] {
        guard !routes.isEmpty else { return [] }

        // Group routes by primary city (the city with the most games in the route)
        var routesByRegion: [String: [[Game]]] = [:]

        for route in routes {
            let primaryCity = getPrimaryCity(for: route, stadiums: stadiums)
            routesByRegion[primaryCity, default: []].append(route)
        }

        // Sort routes within each region by score (best first)
        for (region, regionRoutes) in routesByRegion {
            routesByRegion[region] = regionRoutes.sorted {
                scorePath($0, stadiums: stadiums) > scorePath($1, stadiums: stadiums)
            }
        }

        // Sort regions by their best route's score (so best regions come first)
        let sortedRegions = routesByRegion.keys.sorted { region1, region2 in
            let score1 = routesByRegion[region1]?.first.map { scorePath($0, stadiums: stadiums) } ?? 0
            let score2 = routesByRegion[region2]?.first.map { scorePath($0, stadiums: stadiums) } ?? 0
            return score1 > score2
        }

        print("[GameDAGRouter] Found \(sortedRegions.count) distinct regions: \(sortedRegions.prefix(10).joined(separator: ", "))")

        // Pick routes round-robin from each region to ensure diversity
        var selectedRoutes: [[Game]] = []
        var regionIndices: [String: Int] = [:]

        // First pass: get best route from each region
        for region in sortedRegions {
            if selectedRoutes.count >= maxCount { break }
            if let regionRoutes = routesByRegion[region], !regionRoutes.isEmpty {
                selectedRoutes.append(regionRoutes[0])
                regionIndices[region] = 1
            }
        }

        // Second pass: fill remaining slots with next-best routes from top regions
        var round = 1
        while selectedRoutes.count < maxCount {
            var addedAny = false
            for region in sortedRegions {
                if selectedRoutes.count >= maxCount { break }
                let idx = regionIndices[region] ?? 0
                if let regionRoutes = routesByRegion[region], idx < regionRoutes.count {
                    selectedRoutes.append(regionRoutes[idx])
                    regionIndices[region] = idx + 1
                    addedAny = true
                }
            }
            if !addedAny { break }
            round += 1
            if round > 5 { break } // Safety limit
        }

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

        // Sort by score (best first)
        let sorted = uniquePaths.sorted { scorePath($0, stadiums: stadiums) > scorePath($1, stadiums: stadiums) }

        // Dominance pruning: within same ending city, keep only best paths
        var pruned: [[Game]] = []
        var bestByEndCity: [String: Double] = [:]

        for path in sorted {
            guard let lastGame = path.last else { continue }
            let endCity = stadiums[lastGame.stadiumId]?.city ?? "Unknown"
            let score = scorePath(path, stadiums: stadiums)

            // Keep if this is the best path ending in this city, or if score is within 20% of best
            if let bestScore = bestByEndCity[endCity] {
                if score >= bestScore * 0.8 {
                    pruned.append(path)
                }
            } else {
                bestByEndCity[endCity] = score
                pruned.append(path)
            }

            // Stop if we have enough
            if pruned.count >= beamWidth * 2 {
                break
            }
        }

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