backend/optimizer.py

import random
import math
from datetime import datetime
from backend.routing import RoutingService

class PatrolOptimizer:
    def __init__(self, data_engine, routing_service: RoutingService, geometry_router: RoutingService = None):
        self.data_engine = data_engine
        self.router = routing_service
        self.geometry_router = geometry_router or routing_service

    def _get_station_centroid(self, police_station: str) -> tuple:
        """Calculate the average coordinate of violations for a station to act as the station HQ."""
        lats = []
        lons = []
        for r in self.data_engine.records:
            if r['police_station'] == police_station:
                lats.append(r['lat'])
                lons.append(r['lon'])
                
        if lats:
            return (sum(lats) / len(lats), sum(lons) / len(lons))
        # Default Bengaluru center if police station is not found
        return (12.9716, 77.5946)

    def optimize_routes(self, 
                        police_station: str, 
                        day_of_week: int, 
                        hour: int, 
                        num_vehicles: int = 2, 
                        time_budget_min: float = 120.0,
                        clearance_time_min: float = 15.0,
                        use_predictions: bool = True,
                        predictor = None,
                        strategy: str = "balanced",
                        custom_hotspots: list = None,
                        cleared_hotspots: set = None) -> dict:
        """
        Generate optimal routes for multiple patrol vehicles.
        
        Args:
            police_station (str): Station name (e.g. 'Upparpet')
            day_of_week (int): 0-6
            hour (int): 0-23
            num_vehicles (int): Number of separate patrol routes to generate
            time_budget_min (float): Shift duration in minutes
            clearance_time_min (float): Time spent at each hotspot in minutes
            use_predictions (bool): Whether to use ML predicted GCI or historical GCI
            predictor (CongestionPredictor): Pre-loaded predictor instance
            strategy (str): Routing strategy
            custom_hotspots (list): List of custom manually reported hotspots
            cleared_hotspots (set): Set of manually cleared hotspot IDs or grid cell keys
            
        Returns:
            dict: Structured routes data for visual rendering
        """
        start_time = datetime.now()
        
        # 1. Fetch grid congestion indices
        if use_predictions and predictor is not None:
            grid_data = predictor.predict_grid_gci(day_of_week, hour, police_station)
        else:
            grid_data = self.data_engine.compute_grid_gci(day_of_week, hour, police_station)
            
        # Filter out cleared historical grid cells
        if cleared_hotspots:
            for key in list(grid_data.keys()):
                if key in cleared_hotspots:
                    grid_data[key]['gci'] = 0.0

        # 2. Get police station starting coordinates
        hq_coord = self._get_station_centroid(police_station)
        
        # 3. Filter and sort hotspots
        # Convert grid data into a list of hotspots
        hotspots = []
        for key, info in grid_data.items():
            if info['gci'] > 0.0:
                hotspots.append({
                    'id': key,
                    'x': info['x'],
                    'y': info['y'],
                    'lat': info['lat'],
                    'lon': info['lon'],
                    'gci': info['gci'],
                    'count': info['count'],
                    'explanation': info.get('explanation', {})
                })
                
        # Append active custom hotspots
        if custom_hotspots:
            for ch in custom_hotspots:
                ch_id = ch['id']
                if cleared_hotspots and (ch_id in cleared_hotspots or f"{ch['x']}_{ch['y']}" in cleared_hotspots):
                    continue
                hotspots.append({
                    'id': ch_id,
                    'x': ch['x'],
                    'y': ch['y'],
                    'lat': ch['lat'],
                    'lon': ch['lon'],
                    'gci': ch['gci'],
                    'count': ch.get('count', 1),
                    'explanation': ch.get('explanation', {
                        'priority_tier': ch.get('severity', 'High'),
                        'top_vehicle_type': ch.get('vehicle_type', 'Car'),
                        'top_violation_type': ch.get('violation_type', 'Obstruction'),
                        'peak_hour': hour,
                        'confidence': 'High (Manual Dispatch)',
                        'reason': f"Live Report: {ch.get('description', '')}"
                    })
                })

        # Sort by congestion index descending
        strategy = strategy if strategy in {"max_impact", "fast_sweep", "balanced"} else "balanced"
        hotspots = self._rank_hotspots_for_strategy(hotspots, hq_coord, strategy)
        
        # Limit to top 40 hotspots to keep routing computations fast and relevant
        hotspots = hotspots[:40]
        
        routes = []
        visited_ids = set()
        
        # 4. Generate routes sequentially for each vehicle
        for v_idx in range(num_vehicles):
            # Exclude already visited hotspots
            available_hotspots = [h for h in hotspots if h['id'] not in visited_ids]
            
            if not available_hotspots:
                break
                
            # Run Genetic Algorithm for this vehicle
            route_path, route_stats = self._run_genetic_algorithm(
                hq_coord, 
                available_hotspots, 
                time_budget_min, 
                clearance_time_min,
                strategy
            )
            
            # Record visited nodes to avoid duplicates across vehicles
            for node in route_path:
                visited_ids.add(node['id'])
                
            road_geometry = self._build_route_geometry(hq_coord, route_path)

            routes.append({
                'vehicle_id': v_idx + 1,
                'path': route_path,
                'road_path': road_geometry['road_path'],
                'legs': road_geometry['legs'],
                'route_provider': road_geometry['provider'],
                'is_road_geometry': road_geometry['is_road_geometry'],
                'stats': route_stats
            })
            
        duration_ms = (datetime.now() - start_time).total_seconds() * 1000.0
        impact_summary = self._build_impact_summary(hotspots, routes, time_budget_min)
        
        return {
            'police_station': police_station,
            'strategy': strategy,
            'hq_coordinate': hq_coord,
            'routes': routes,
            'impact_summary': impact_summary,
            'computation_time_ms': round(duration_ms, 1)
        }

    def _build_route_geometry(self, hq_coord: tuple, route_path: list) -> dict:
        if not route_path:
            return {
                'road_path': [],
                'legs': [],
                'provider': 'none',
                'is_road_geometry': False
            }

        legs = []
        road_path = []
        provider = 'none'
        is_road_geometry = False
        current_coord = hq_coord

        for node in route_path:
            next_coord = (node['lat'], node['lon'])
            leg = self.geometry_router.get_route(current_coord, next_coord)
            provider = leg.get('provider', provider)
            is_road_geometry = is_road_geometry or leg.get('is_road_geometry', False)
            coords = leg.get('path_coords', [])

            if not road_path:
                road_path.extend(coords)
            elif coords:
                road_path.extend(coords[1:])

            legs.append({
                'from': [round(current_coord[0], 6), round(current_coord[1], 6)],
                'to': [round(next_coord[0], 6), round(next_coord[1], 6)],
                'distance_km': leg.get('distance_km', 0),
                'duration_min': leg.get('duration_min', 0),
                'provider': leg.get('provider', 'unknown'),
                'is_road_geometry': leg.get('is_road_geometry', False),
                'path_coords': coords
            })
            current_coord = next_coord

        return {
            'road_path': road_path if is_road_geometry else [],
            'legs': legs,
            'provider': provider,
            'is_road_geometry': is_road_geometry
        }

    def _rank_hotspots_for_strategy(self, hotspots: list, hq_coord: tuple, strategy: str) -> list:
        if strategy == "fast_sweep":
            return sorted(hotspots, key=lambda x: (x.get('count', 0), x['gci']), reverse=True)
        if strategy == "max_impact":
            return sorted(hotspots, key=lambda x: x['gci'], reverse=True)

        def balanced_score(h):
            route = self.router.get_route(hq_coord, (h['lat'], h['lon']))
            return h['gci'] / max(1.0, route['duration_min'])

        return sorted(hotspots, key=balanced_score, reverse=True)

    def _build_impact_summary(self, hotspots: list, routes: list, time_budget_min: float) -> dict:
        total_hotspot_gci = sum(h['gci'] for h in hotspots)
        cleared_ids = set()
        cleared_gci = 0.0
        cleared_count = 0
        total_time = 0.0

        for route in routes:
            total_time += route['stats']['total_time_min']
            for node in route['path']:
                if node['id'] not in cleared_ids:
                    cleared_ids.add(node['id'])
                    cleared_gci += node['gci']
                    cleared_count += node.get('count', 0)

        remaining_gci = max(0.0, total_hotspot_gci - cleared_gci)
        reduction_pct = (cleared_gci / total_hotspot_gci * 100.0) if total_hotspot_gci else 0.0
        avg_utilization = total_time / max(1.0, len(routes) * time_budget_min) * 100.0 if routes else 0.0

        if reduction_pct >= 55.0:
            note = "High-impact deployment: patrols clear most priority hotspot load in this window."
        elif reduction_pct >= 30.0:
            note = "Balanced deployment: meaningful hotspot relief while preserving patrol coverage."
        elif reduction_pct > 0.0:
            note = "Limited relief: increase units or time budget for stronger congestion reduction."
        else:
            note = "No clearable hotspots found under the selected constraints."

        return {
            'before_gci': round(total_hotspot_gci, 1),
            'after_gci': round(remaining_gci, 1),
            'cleared_gci': round(cleared_gci, 1),
            'reduction_pct': round(reduction_pct, 1),
            'cleared_hotspots': len(cleared_ids),
            'estimated_records_impacted': cleared_count,
            'avg_unit_utilization_pct': round(avg_utilization, 1),
            'cleared_grid_ids': sorted(cleared_ids),
            'recommendation': note
        }

    def _run_genetic_algorithm(self, 
                               hq_coord: tuple, 
                               hotspots: list, 
                               budget: float, 
                               clearance: float,
                               strategy: str) -> tuple:
        """Solves the Orienteering Problem using a Genetic Algorithm."""
        if not hotspots:
            return [], {'total_gci': 0, 'total_time_min': 0, 'hotspots_visited': 0}
            
        # GA Hyperparameters
        pop_size = 40
        generations = 60
        mutation_rate = 0.15
        
        # Precompute distance/time from HQ to all nodes, and between all nodes to save compute cycles
        # cache key: (coordA, coordB) -> (distance, time_min)
        dist_cache = {}
        
        def get_travel_time(coord1, coord2):
            cache_key = (coord1, coord2)
            if cache_key not in dist_cache:
                route = self.router.get_route(coord1, coord2)
                dist_cache[cache_key] = route['duration_min']
            return dist_cache[cache_key]
            
        # Helper: Decode chromosome (a permutation of indices) to a valid route
        def decode(chromosome):
            path = []
            total_time = 0.0
            total_gci = 0.0
            total_count = 0
            
            current_coord = hq_coord
            
            for idx in chromosome:
                h = hotspots[idx]
                h_coord = (h['lat'], h['lon'])
                
                # Calculate time to travel from current location to this hotspot
                travel_time = get_travel_time(current_coord, h_coord)
                
                # Check if visiting this hotspot fits in time budget
                if total_time + travel_time + clearance <= budget:
                    total_time += travel_time + clearance
                    total_gci += h['gci']
                    total_count += h.get('count', 0)
                    path.append(h)
                    current_coord = h_coord
                else:
                    # Once budget is exceeded, stop adding nodes (truncation)
                    break
                    
            if strategy == "fast_sweep":
                fitness = len(path) * 35.0 + total_count * 2.0 - total_time * 0.15
            elif strategy == "max_impact":
                fitness = total_gci
            else:
                fitness = total_gci + len(path) * 8.0 - total_time * 0.35

            return path, total_gci, total_time, fitness

        # Initialize Population (random permutations of hotspot indices)
        num_nodes = len(hotspots)
        population = []
        for _ in range(pop_size):
            chrom = list(range(num_nodes))
            random.shuffle(chrom)
            population.append(chrom)
            
        # GA Loop
        for _ in range(generations):
            # 1. Evaluate fitness
            fitness_scores = []
            for chrom in population:
                _, _, _, fitness = decode(chrom)
                fitness_scores.append(fitness)
                
            # 2. Selection (Tournament Selection)
            new_pop = []
            for _ in range(pop_size):
                # Sample 3 random individuals and select the best
                candidates = random.sample(list(zip(population, fitness_scores)), 3)
                best = max(candidates, key=lambda x: x[1])[0]
                new_pop.append(best.copy())
                
            # 3. Crossover (Ordered Crossover for permutations)
            for i in range(0, pop_size, 2):
                if i + 1 < pop_size and random.random() < 0.8:
                    p1, p2 = new_pop[i], new_pop[i+1]
                    
                    # Crossover slice
                    cx1 = random.randint(0, num_nodes - 2)
                    cx2 = random.randint(cx1 + 1, num_nodes - 1)
                    
                    # Child 1
                    c1 = [-1] * num_nodes
                    c1[cx1:cx2] = p1[cx1:cx2]
                    # Fill remainder from p2
                    p2_filtered = [item for item in p2 if item not in c1[cx1:cx2]]
                    c1[:cx1] = p2_filtered[:cx1]
                    c1[cx2:] = p2_filtered[cx1:]
                    
                    # Child 2
                    c2 = [-1] * num_nodes
                    c2[cx1:cx2] = p2[cx1:cx2]
                    # Fill remainder from p1
                    p1_filtered = [item for item in p1 if item not in c2[cx1:cx2]]
                    c2[:cx1] = p1_filtered[:cx1]
                    c2[cx2:] = p1_filtered[cx1:]
                    
                    new_pop[i], new_pop[i+1] = c1, c2
                    
            # 4. Mutation (Swap Mutation)
            for i in range(pop_size):
                if random.random() < mutation_rate:
                    m1 = random.randint(0, num_nodes - 1)
                    m2 = random.randint(0, num_nodes - 1)
                    new_pop[i][m1], new_pop[i][m2] = new_pop[i][m2], new_pop[i][m1]
                    
            population = new_pop

        # Get the absolute best solution from final population
        best_chrom = max(population, key=lambda chrom: decode(chrom)[3])
        best_path, best_gci, best_time, _ = decode(best_chrom)
        
        # Format the stats
        stats = {
            'total_gci': round(best_gci, 1),
            'total_time_min': round(best_time, 1),
            'hotspots_visited': len(best_path)
        }
        
        return best_path, stats