"""
Route optimization for cave surveys.
Finds optimal paths through cave systems considering stability and safety.

Author: Samir Baladi
"""

import numpy as np
from typing import Dict, List, Tuple, Optional, Union
from dataclasses import dataclass
import heapq
import logging

logger = logging.getLogger(__name__)

@dataclass
class RouteNode:
    """Represents a node in the route network."""
    id: str
    x: float
    y: float
    z: float
    stability: float  # 0-1 stability score
    connections: List[Tuple[str, float]]  # (node_id, distance)


class RouteOptimization:
    """
    Optimizes survey routes through cave systems.
    Uses A* algorithm with stability weighting.
    """
    
    def __init__(self):
        """Initialize route optimization."""
        self.nodes = {}
        self.routes = []
        logger.info("Initialized RouteOptimization")

    def add_node(self,
                node_id: str,
                x: float,
                y: float,
                z: float,
                stability: float = 1.0) -> None:
        """
        Add a node to the route network.
        
        Args:
            node_id: Node identifier
            x, y, z: Coordinates
            stability: Node stability score (0-1)
        """
        node = RouteNode(
            id=node_id,
            x=x, y=y, z=z,
            stability=stability,
            connections=[]
        )
        
        self.nodes[node_id] = node
        logger.debug(f"Added node: {node_id}")

    def add_connection(self,
                      node1_id: str,
                      node2_id: str,
                      distance: Optional[float] = None) -> None:
        """
        Add a connection between nodes.
        
        Args:
            node1_id: First node
            node2_id: Second node
            distance: Distance (calculated if None)
        """
        if node1_id not in self.nodes or node2_id not in self.nodes:
            logger.warning(f"Cannot connect {node1_id} - {node2_id}: nodes not found")
            return
        
        node1 = self.nodes[node1_id]
        node2 = self.nodes[node2_id]
        
        if distance is None:
            # Calculate Euclidean distance
            distance = np.sqrt(
                (node1.x - node2.x)**2 +
                (node1.y - node2.y)**2 +
                (node1.z - node2.z)**2
            )
        
        # Add bidirectional connection
        node1.connections.append((node2_id, distance))
        node2.connections.append((node1_id, distance))
        
        logger.debug(f"Added connection: {node1_id} - {node2_id} ({distance:.1f}m)")

    def find_optimal_route(self,
                          start_id: str,
                          end_id: str,
                          safety_weight: float = 0.5) -> Dict[str, Union[List[str], float]]:
        """
        Find optimal route between two nodes.
        
        Uses A* algorithm with cost function:
        cost = distance * (1 - safety_weight * (1 - stability))
        
        Args:
            start_id: Start node
            end_id: End node
            safety_weight: Weight of safety vs distance (0-1)
            
        Returns:
            Dictionary with route information
        """
        if start_id not in self.nodes or end_id not in self.nodes:
            return {'error': 'Start or end node not found'}
        
        # A* algorithm
        open_set = [(0, start_id)]
        came_from = {}
        g_score = {node_id: float('inf') for node_id in self.nodes}
        g_score[start_id] = 0
        f_score = {node_id: float('inf') for node_id in self.nodes}
        f_score[start_id] = self._heuristic(start_id, end_id)
        
        open_set_hash = {start_id}
        
        while open_set:
            current = heapq.heappop(open_set)[1]
            open_set_hash.remove(current)
            
            if current == end_id:
                # Reconstruct path
                path = []
                while current in came_from:
                    path.append(current)
                    current = came_from[current]
                path.append(start_id)
                path.reverse()
                
                # Calculate metrics
                total_distance = g_score[end_id]
                avg_stability = self._calculate_path_stability(path)
                
                route = {
                    'path': path,
                    'total_distance': round(total_distance, 1),
                    'avg_stability': round(avg_stability, 2),
                    'n_nodes': len(path),
                    'safety_weight': safety_weight
                }
                
                self.routes.append(route)
                return route
            
            for neighbor_id, distance in self.nodes[current].connections:
                # Calculate cost with safety weighting
                neighbor_stability = self.nodes[neighbor_id].stability
                safety_cost = distance * (1 - safety_weight * (1 - neighbor_stability))
                
                tentative_g = g_score[current] + safety_cost
                
                if tentative_g < g_score[neighbor_id]:
                    came_from[neighbor_id] = current
                    g_score[neighbor_id] = tentative_g
                    f_score[neighbor_id] = tentative_g + self._heuristic(neighbor_id, end_id)
                    
                    if neighbor_id not in open_set_hash:
                        heapq.heappush(open_set, (f_score[neighbor_id], neighbor_id))
                        open_set_hash.add(neighbor_id)
        
        return {'error': 'No path found'}

    def _heuristic(self, node_id: str, target_id: str) -> float:
        """Heuristic for A* (Euclidean distance)."""
        node = self.nodes[node_id]
        target = self.nodes[target_id]
        
        return np.sqrt(
            (node.x - target.x)**2 +
            (node.y - target.y)**2 +
            (node.z - target.z)**2
        )

    def _calculate_path_stability(self, path: List[str]) -> float:
        """Calculate average stability along path."""
        stabilities = [self.nodes[node_id].stability for node_id in path]
        return np.mean(stabilities)

    def find_safest_route(self,
                         start_id: str,
                         end_id: str,
                         min_stability: float = 0.7) -> Dict[str, Union[List[str], float]]:
        """
        Find safest route (maximizes minimum stability).
        
        Args:
            start_id: Start node
            end_id: End node
            min_stability: Minimum acceptable stability
            
        Returns:
            Dictionary with route information
        """
        # Filter nodes by stability
        valid_nodes = {nid for nid, node in self.nodes.items() 
                      if node.stability >= min_stability}
        
        if start_id not in valid_nodes or end_id not in valid_nodes:
            return {'error': 'Start or end node below stability threshold'}
        
        # Build subgraph with valid nodes
        from collections import defaultdict
        subgraph = defaultdict(list)
        
        for node_id in valid_nodes:
            node = self.nodes[node_id]
            for neighbor_id, distance in node.connections:
                if neighbor_id in valid_nodes:
                    subgraph[node_id].append((neighbor_id, distance))
        
        # BFS for path maximizing minimum stability
        from collections import deque
        
        queue = deque([(start_id, [start_id], float('inf'))])
        best_path = None
        best_min_stability = 0
        
        while queue:
            current, path, min_stab = queue.popleft()
            
            if current == end_id:
                if min_stab > best_min_stability:
                    best_min_stability = min_stab
                    best_path = path
                continue
            
            for neighbor, _ in subgraph[current]:
                if neighbor not in path:
                    new_min_stab = min(min_stab, self.nodes[neighbor].stability)
                    queue.append((neighbor, path + [neighbor], new_min_stab))
        
        if best_path:
            # Calculate distance
            distance = 0
            for i in range(len(best_path) - 1):
                for conn_id, dist in self.nodes[best_path[i]].connections:
                    if conn_id == best_path[i + 1]:
                        distance += dist
                        break
            
            return {
                'path': best_path,
                'total_distance': round(distance, 1),
                'min_stability': round(best_min_stability, 2),
                'avg_stability': round(self._calculate_path_stability(best_path), 2),
                'n_nodes': len(best_path)
            }
        
        return {'error': 'No safe path found'}

    def calculate_route_risk(self,
                            path: List[str],
                            hazard_zones: List[Dict]) -> Dict[str, Union[str, float]]:
        """
        Calculate risk for a given route.
        
        Args:
            path: List of node IDs
            hazard_zones: List of hazard zones with positions and risk levels
            
        Returns:
            Dictionary with risk assessment
        """
        total_risk = 0
        max_risk = 0
        risk_by_node = []
        
        for node_id in path:
            node = self.nodes[node_id]
            node_risk = 0
            
            # Check proximity to hazard zones
            for zone in hazard_zones:
                zone_pos = zone['position']
                distance = np.sqrt(
                    (node.x - zone_pos[0])**2 +
                    (node.y - zone_pos[1])**2 +
                    (node.z - zone_pos[2])**2
                )
                
                if distance < zone.get('radius', 50):
                    risk_factor = 1 - distance / zone['radius']
                    zone_risk = risk_factor * zone.get('risk_level', 1)
                    node_risk = max(node_risk, zone_risk)
            
            risk_by_node.append(node_risk)
            total_risk += node_risk
            max_risk = max(max_risk, node_risk)
        
        avg_risk = total_risk / len(path) if path else 0
        
        if max_risk > 0.8:
            assessment = "EXTREME RISK"
            color = "⚫"
        elif max_risk > 0.6:
            assessment = "HIGH RISK"
            color = "🔴"
        elif max_risk > 0.4:
            assessment = "MODERATE RISK"
            color = "🟠"
        elif max_risk > 0.2:
            assessment = "LOW RISK"
            color = "🟡"
        else:
            assessment = "MINIMAL RISK"
            color = "🟢"
        
        return {
            'assessment': assessment,
            'color': color,
            'max_risk': round(max_risk, 2),
            'avg_risk': round(avg_risk, 2),
            'risk_by_node': [round(r, 2) for r in risk_by_node]
        }

    def optimize_survey_order(self,
                             nodes_to_visit: List[str],
                             start_node: str) -> Dict[str, Union[List[str], float]]:
        """
        Optimize order of visiting multiple nodes (traveling salesman problem).
        
        Args:
            nodes_to_visit: List of nodes to visit
            start_node: Starting node
            
        Returns:
            Dictionary with optimized route
        """
        if start_node not in self.nodes:
            return {'error': 'Start node not found'}
        
        # Ensure all nodes exist
        for node_id in nodes_to_visit:
            if node_id not in self.nodes:
                return {'error': f'Node {node_id} not found'}
        
        # Simple nearest neighbor heuristic
        unvisited = set(nodes_to_visit)
        current = start_node
        route = [current]
        total_distance = 0
        
        while unvisited:
            # Find nearest unvisited node
            nearest = None
            nearest_dist = float('inf')
            
            for node_id in unvisited:
                # Find shortest path to this node
                path_result = self.find_optimal_route(current, node_id, safety_weight=0)
                if 'error' not in path_result:
                    dist = path_result['total_distance']
                    if dist < nearest_dist:
                        nearest_dist = dist
                        nearest = node_id
            
            if nearest is None:
                break
            
            route.append(nearest)
            total_distance += nearest_dist
            unvisited.remove(nearest)
            current = nearest
        
        return {
            'route': route,
            'total_distance': round(total_distance, 1),
            'n_nodes': len(route)
        }

    def summary(self) -> Dict[str, Union[int, float, str]]:
        """Get summary of route optimization."""
        return {
            'n_nodes': len(self.nodes),
            'n_connections': sum(len(node.connections) for node in self.nodes.values()) // 2,
            'n_routes_found': len(self.routes),
            'avg_node_stability': round(np.mean([n.stability for n in self.nodes.values()]), 2)
        }


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Route Optimization")
    print("="*60)
    
    ro = RouteOptimization()
    
    # Add nodes
    ro.add_node("Entrance", 0, 0, 0, stability=1.0)
    ro.add_node("Junction", 100, 50, -50, stability=0.9)
    ro.add_node("Chamber", 200, 100, -120, stability=0.7)
    ro.add_node("Sump", 250, 150, -150, stability=0.5)
    
    # Add connections
    ro.add_connection("Entrance", "Junction")
    ro.add_connection("Junction", "Chamber")
    ro.add_connection("Chamber", "Sump")
    ro.add_connection("Entrance", "Chamber")  # Direct connection
    
    # Find optimal route
    route = ro.find_optimal_route("Entrance", "Sump", safety_weight=0.7)
    if 'path' in route:
        print(f"\nOptimal Route:")
        print(f"  Path: {' → '.join(route['path'])}")
        print(f"  Distance: {route['total_distance']}m")
        print(f"  Avg Stability: {route['avg_stability']}")
    
    # Find safest route
    safest = ro.find_safest_route("Entrance", "Sump", min_stability=0.6)
    if 'path' in safest:
        print(f"\nSafest Route:")
        print(f"  Path: {' → '.join(safest['path'])}")
        print(f"  Min Stability: {safest['min_stability']}")
    
    print("\n" + "="*60)
