"""
Loop closure analysis for cave surveys.
Calculates and analyzes closure errors in survey loops.

Author: Samir Baladi
"""

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

logger = logging.getLogger(__name__)

@dataclass
class Loop:
    """Represents a survey loop."""
    stations: List[str]
    length: float  # Total loop length (m)
    error_xy: float  # Horizontal closure error (m)
    error_z: float  # Vertical closure error (m)
    error_3d: float  # 3D closure error (m)
    error_percent: float  # Error as percentage of length


class ClosureError:
    """
    Calculates and analyzes closure errors in survey loops.
    """
    
    def __init__(self, max_acceptable_error: float = 0.5):
        """
        Initialize closure error analyzer.
        
        Args:
            max_acceptable_error: Maximum acceptable error as percentage of loop length
        """
        self.max_acceptable_error = max_acceptable_error
        self.loops = []
        self.stations = {}
        logger.info(f"Initialized ClosureError with max error {max_acceptable_error}%")

    def add_station(self, name: str, x: float, y: float, z: float) -> None:
        """
        Add a station coordinate.
        
        Args:
            name: Station name
            x, y, z: Coordinates
        """
        self.stations[name] = (x, y, z)
        logger.debug(f"Added station: {name}")

    def find_loops(self, connections: List[Tuple[str, str]], max_loop_length: int = 10) -> List[Loop]:
        """
        Find all possible survey loops.
        
        Args:
            connections: List of (from_station, to_station) connections
            max_loop_length: Maximum number of stations in loop
            
        Returns:
            List of loops found
        """
        # Build adjacency list
        adj = {}
        for a, b in connections:
            if a not in adj:
                adj[a] = []
            if b not in adj:
                adj[b] = []
            adj[a].append(b)
            adj[b].append(a)
        
        # Find cycles using DFS
        loops = []
        
        for start in adj.keys():
            stack = [(start, [start])]
            
            while stack:
                node, path = stack.pop()
                
                for neighbor in adj[node]:
                    if neighbor == start and len(path) > 2:
                        # Found a cycle
                        if len(path) <= max_loop_length:
                            loops.append(path)
                    elif neighbor not in path and len(path) < max_loop_length:
                        stack.append((neighbor, path + [neighbor]))
        
        # Remove duplicates (cycles with different starting points)
        unique_loops = []
        loop_strings = set()
        
        for loop in loops:
            # Normalize loop representation
            min_index = loop.index(min(loop))
            normalized = loop[min_index:] + loop[:min_index]
            loop_str = '->'.join(normalized)
            
            if loop_str not in loop_strings:
                loop_strings.add(loop_str)
                unique_loops.append(loop)
        
        # Calculate errors for each loop
        for loop_stations in unique_loops:
            error = self.calculate_loop_error(loop_stations)
            if error['error_3d'] > 0.01:  # Only report significant errors
                loop = Loop(
                    stations=loop_stations,
                    length=error['length'],
                    error_xy=error['error_xy'],
                    error_z=error['error_z'],
                    error_3d=error['error_3d'],
                    error_percent=error['error_percent']
                )
                self.loops.append(loop)
        
        return self.loops

    def calculate_loop_error(self, stations: List[str]) -> Dict[str, float]:
        """
        Calculate closure error for a specific loop.
        
        Args:
            stations: List of station names in the loop
            
        Returns:
            Dictionary with error metrics
        """
        if len(stations) < 3:
            return {'error_3d': 0, 'length': 0}
        
        # Calculate cumulative vectors
        x_sum, y_sum, z_sum = 0, 0, 0
        length = 0
        
        for i in range(len(stations)):
            from_station = stations[i]
            to_station = stations[(i + 1) % len(stations)]
            
            if from_station in self.stations and to_station in self.stations:
                x1, y1, z1 = self.stations[from_station]
                x2, y2, z2 = self.stations[to_station]
                
                dx = x2 - x1
                dy = y2 - y1
                dz = z2 - z1
                
                x_sum += dx
                y_sum += dy
                z_sum += dz
                
                length += np.sqrt(dx**2 + dy**2 + dz**2)
        
        # Calculate errors
        error_xy = np.sqrt(x_sum**2 + y_sum**2)
        error_z = abs(z_sum)
        error_3d = np.sqrt(x_sum**2 + y_sum**2 + z_sum**2)
        
        # Error as percentage of total length
        if length > 0:
            error_percent = (error_3d / length) * 100
        else:
            error_percent = 0
        
        return {
            'length': round(length, 2),
            'error_xy': round(error_xy, 3),
            'error_z': round(error_z, 3),
            'error_3d': round(error_3d, 3),
            'error_percent': round(error_percent, 3)
        }

    def adjust_loop(self, stations: List[str]) -> Dict[str, Dict[str, float]]:
        """
        Adjust coordinates to close a loop.
        
        Args:
            stations: List of station names in the loop
            
        Returns:
            Dictionary with adjusted coordinates
        """
        if len(stations) < 3:
            return {}
        
        # Calculate closure error
        error = self.calculate_loop_error(stations)
        
        # Distribute error proportionally along loop
        adjusted = {}
        
        # Simplified adjustment - distribute error evenly
        n_segments = len(stations)
        correction_per_segment = {
            'x': -error['error_xy'] * 0.5 / n_segments,
            'y': -error['error_xy'] * 0.5 / n_segments,
            'z': -error['error_z'] / n_segments
        }
        
        cumulative_correction = {'x': 0, 'y': 0, 'z': 0}
        
        for i, station in enumerate(stations):
            if station in self.stations:
                x, y, z = self.stations[station]
                
                # Apply cumulative correction
                adj_x = x + cumulative_correction['x']
                adj_y = y + cumulative_correction['y']
                adj_z = z + cumulative_correction['z']
                
                adjusted[station] = {
                    'original': (x, y, z),
                    'adjusted': (adj_x, adj_y, adj_z),
                    'correction': (cumulative_correction['x'], 
                                  cumulative_correction['y'], 
                                  cumulative_correction['z'])
                }
                
                # Update cumulative correction for next station
                cumulative_correction['x'] += correction_per_segment['x']
                cumulative_correction['y'] += correction_per_segment['y']
                cumulative_correction['z'] += correction_per_segment['z']
        
        return adjusted

    def analyze_error_distribution(self) -> Dict[str, Union[float, str]]:
        """
        Analyze distribution of closure errors.
        
        Returns:
            Dictionary with error statistics
        """
        if not self.loops:
            return {'status': 'No loops found'}
        
        errors_xy = [l.error_xy for l in self.loops]
        errors_z = [l.error_z for l in self.loops]
        errors_3d = [l.error_3d for l in self.loops]
        errors_percent = [l.error_percent for l in self.loops]
        
        # Calculate statistics
        stats = {
            'n_loops': len(self.loops),
            'mean_error_xy': round(np.mean(errors_xy), 3),
            'max_error_xy': round(max(errors_xy), 3),
            'mean_error_z': round(np.mean(errors_z), 3),
            'max_error_z': round(max(errors_z), 3),
            'mean_error_3d': round(np.mean(errors_3d), 3),
            'max_error_3d': round(max(errors_3d), 3),
            'mean_error_percent': round(np.mean(errors_percent), 3),
            'max_error_percent': round(max(errors_percent), 3)
        }
        
        # Grade survey quality
        mean_percent = stats['mean_error_percent']
        if mean_percent < 0.1:
            grade = "Excellent"
            color = "🟢"
        elif mean_percent < 0.3:
            grade = "Good"
            color = "🟢"
        elif mean_percent < 0.5:
            grade = "Fair"
            color = "🟡"
        elif mean_percent < 1.0:
            grade = "Poor"
            color = "🟠"
        else:
            grade = "Unacceptable"
            color = "🔴"
        
        stats['survey_grade'] = grade
        stats['color'] = color
        
        return stats

    def identify_worst_loops(self, n: int = 5) -> List[Dict]:
        """
        Identify loops with worst closure errors.
        
        Args:
            n: Number of loops to return
            
        Returns:
            List of worst loops
        """
        if not self.loops:
            return []
        
        sorted_loops = sorted(self.loops, key=lambda l: l.error_percent, reverse=True)
        
        worst = []
        for loop in sorted_loops[:n]:
            worst.append({
                'stations': ' → '.join(loop.stations),
                'length': round(loop.length, 1),
                'error_3d': round(loop.error_3d, 3),
                'error_percent': round(loop.error_percent, 2),
                'error_xy': round(loop.error_xy, 3),
                'error_z': round(loop.error_z, 3)
            })
        
        return worst

    def calculate_loop_statistics(self) -> Dict[str, float]:
        """
        Calculate comprehensive loop statistics.
        
        Returns:
            Dictionary with loop statistics
        """
        if not self.loops:
            return {}
        
        lengths = [l.length for l in self.loops]
        
        return {
            'total_loops': len(self.loops),
            'total_loop_length': round(sum(lengths), 1),
            'mean_loop_length': round(np.mean(lengths), 1),
            'max_loop_length': round(max(lengths), 1),
            'min_loop_length': round(min(lengths), 1)
        }

    def check_acceptability(self, error_percent: float) -> Dict[str, Union[bool, str]]:
        """
        Check if error is within acceptable limits.
        
        Args:
            error_percent: Closure error as percentage
            
        Returns:
            Dictionary with acceptability assessment
        """
        acceptable = error_percent <= self.max_acceptable_error
        
        if acceptable:
            status = "ACCEPTABLE"
            action = "Continue with survey"
        else:
            status = "EXCESSIVE"
            action = "Re-survey loop or apply adjustment"
        
        return {
            'acceptable': acceptable,
            'status': status,
            'action': action,
            'error': error_percent,
            'limit': self.max_acceptable_error
        }

    def summary(self) -> Dict[str, Union[int, float, str, List]]:
        """Get complete summary of closure error analysis."""
        if not self.loops:
            return {'status': 'No loops found'}
        
        stats = self.analyze_error_distribution()
        worst = self.identify_worst_loops(3)
        loop_stats = self.calculate_loop_statistics()
        
        return {
            'error_statistics': stats,
            'loop_statistics': loop_stats,
            'worst_loops': worst,
            'recommendation': self._get_recommendation(stats['mean_error_percent'])
        }

    def _get_recommendation(self, mean_error: float) -> str:
        """Get recommendation based on error level."""
        if mean_error < 0.1:
            return "Survey quality excellent - proceed normally"
        elif mean_error < 0.3:
            return "Survey quality good - continue monitoring"
        elif mean_error < 0.5:
            return "Moderate errors - check instruments"
        elif mean_error < 1.0:
            return "Significant errors - recalibrate and resurvey critical loops"
        else:
            return "Unacceptable errors - complete resurvey required"


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Closure Error Analysis")
    print("="*60)
    
    ce = ClosureError()
    
    # Add stations
    ce.add_station("A1", 0, 0, 0)
    ce.add_station("A2", 100, 5, -10)
    ce.add_station("A3", 150, 55, -20)
    ce.add_station("A4", 50, 60, -15)
    
    # Define connections
    connections = [
        ("A1", "A2"),
        ("A2", "A3"),
        ("A3", "A4"),
        ("A4", "A1")
    ]
    
    # Find loops
    loops = ce.find_loops(connections)
    print(f"\nFound {len(loops)} loops")
    
    # Analyze error distribution
    stats = ce.analyze_error_distribution()
    print(f"\nError Statistics:")
    print(f"  - Grade: {stats['color']} {stats['survey_grade']}")
    print(f"  - Mean Error: {stats['mean_error_percent']}%")
    print(f"  - Max Error: {stats['max_error_percent']}%")
    
    # Worst loops
    worst = ce.identify_worst_loops(2)
    if worst:
        print(f"\nWorst Loops:")
        for w in worst:
            print(f"  • {w['stations']}: {w['error_percent']}% error")
    
    print("\n" + "="*60)
