"""
Survey station analysis tools.
Analyzes station distribution, connectivity, and survey quality.

Author: Samir Baladi
"""

import numpy as np
from typing import Dict, List, Tuple, Optional, Union
from dataclasses import dataclass
from scipy.spatial import KDTree
import logging
import networkx as nx

logger = logging.getLogger(__name__)

@dataclass
class Station:
    """Represents a survey station with metadata."""
    name: str
    x: float
    y: float
    z: float
    quality: float  # 0-1, based on measurement precision
    connections: List[str]  # Connected stations


class StationAnalysis:
    """
    Analyzes survey station data for quality and network properties.
    """
    
    def __init__(self):
        """Initialize station analysis."""
        self.stations = {}
        self.graph = nx.Graph()
        logger.info("Initialized StationAnalysis")

    def add_station(self,
                   name: str,
                   x: float,
                   y: float,
                   z: float,
                   quality: float = 1.0) -> None:
        """
        Add a survey station.
        
        Args:
            name: Station name
            x, y, z: Coordinates
            quality: Measurement quality (0-1)
        """
        station = Station(
            name=name,
            x=x, y=y, z=z,
            quality=quality,
            connections=[]
        )
        
        self.stations[name] = station
        self.graph.add_node(name, x=x, y=y, z=z, quality=quality)
        logger.debug(f"Added station: {name}")

    def add_connection(self, from_station: str, to_station: str, length: float) -> None:
        """
        Add a connection between stations.
        
        Args:
            from_station: Source station
            to_station: Target station
            length: Shot length
        """
        if from_station in self.stations and to_station in self.stations:
            self.stations[from_station].connections.append(to_station)
            self.stations[to_station].connections.append(from_station)
            self.graph.add_edge(from_station, to_station, weight=length)
            logger.debug(f"Added connection: {from_station} - {to_station}")

    def analyze_network(self) -> Dict[str, Union[float, int, str]]:
        """
        Analyze survey network properties.
        
        Returns:
            Dictionary with network metrics
        """
        if len(self.graph) < 2:
            return {'status': 'Insufficient stations'}
        
        # Basic metrics
        n_nodes = self.graph.number_of_nodes()
        n_edges = self.graph.number_of_edges()
        
        # Connectivity
        if nx.is_connected(self.graph):
            connectivity = "Fully connected"
        else:
            n_components = nx.number_connected_components(self.graph)
            connectivity = f"{n_components} components"
        
        # Degree distribution
        degrees = [d for n, d in self.graph.degree()]
        
        # Centrality
        if n_nodes > 2:
            try:
                betweenness = nx.betweenness_centrality(self.graph)
                most_central = max(betweenness.items(), key=lambda x: x[1])
            except:
                most_central = ("None", 0)
        else:
            most_central = ("None", 0)
        
        # Shortest paths
        if n_nodes > 1 and nx.is_connected(self.graph):
            avg_path_length = nx.average_shortest_path_length(self.graph)
            diameter = nx.diameter(self.graph)
        else:
            avg_path_length = 0
            diameter = 0
        
        return {
            'n_stations': n_nodes,
            'n_shots': n_edges,
            'connectivity': connectivity,
            'mean_degree': round(np.mean(degrees), 2),
            'max_degree': max(degrees),
            'min_degree': min(degrees),
            'most_central_station': most_central[0],
            'centrality_score': round(most_central[1], 3),
            'avg_path_length': round(avg_path_length, 2),
            'diameter': diameter,
            'density': round(nx.density(self.graph), 3)
        }

    def calculate_closure_errors(self) -> List[Dict[str, Union[str, float]]]:
        """
        Calculate closure errors for all loops in the network.
        
        Returns:
            List of closure errors
        """
        closure_errors = []
        
        # Find cycles in the graph
        cycles = list(nx.simple_cycles(self.graph.to_directed()))
        
        for cycle in cycles[:10]:  # Limit to first 10 cycles
            if len(cycle) < 3:
                continue
            
            # Calculate closure error
            x_sum, y_sum, z_sum = 0, 0, 0
            
            for i in range(len(cycle)):
                from_station = cycle[i]
                to_station = cycle[(i + 1) % len(cycle)]
                
                if self.graph.has_edge(from_station, to_station):
                    dx = self.stations[to_station].x - self.stations[from_station].x
                    dy = self.stations[to_station].y - self.stations[from_station].y
                    dz = self.stations[to_station].z - self.stations[from_station].z
                    
                    x_sum += dx
                    y_sum += dy
                    z_sum += dz
            
            error_3d = np.sqrt(x_sum**2 + y_sum**2 + z_sum**2)
            error_xy = np.sqrt(x_sum**2 + y_sum**2)
            
            if error_3d > 0.01:  # Only report significant errors
                closure_errors.append({
                    'loop': ' → '.join(cycle),
                    'length': len(cycle),
                    'error_3d_m': round(error_3d, 3),
                    'error_xy_m': round(error_xy, 3),
                    'error_z_m': round(abs(z_sum), 3),
                    'error_percent': round(error_3d / len(cycle) * 100, 2) if len(cycle) > 0 else 0
                })
        
        return sorted(closure_errors, key=lambda x: x['error_3d_m'], reverse=True)

    def find_isolated_stations(self, max_distance: float = 50) -> List[str]:
        """
        Find isolated stations (far from others).
        
        Args:
            max_distance: Maximum acceptable distance to nearest neighbor (m)
            
        Returns:
            List of isolated station names
        """
        if len(self.stations) < 2:
            return []
        
        # Build coordinate array
        coords = []
        names = []
        for name, station in self.stations.items():
            coords.append([station.x, station.y, station.z])
            names.append(name)
        
        tree = KDTree(coords)
        
        isolated = []
        for i, name in enumerate(names):
            distances, _ = tree.query(coords[i], k=2)  # k=2 to exclude self
            if len(distances) > 1 and distances[1] > max_distance:
                isolated.append(name)
        
        return isolated

    def calculate_station_density(self, volume: Optional[float] = None) -> Dict[str, float]:
        """
        Calculate station density.
        
        Args:
            volume: Survey volume (m³), calculated from extent if None
            
        Returns:
            Dictionary with density metrics
        """
        if len(self.stations) < 2:
            return {'density': 0}
        
        # Calculate survey extent
        xs = [s.x for s in self.stations.values()]
        ys = [s.y for s in self.stations.values()]
        zs = [s.z for s in self.stations.values()]
        
        x_range = max(xs) - min(xs)
        y_range = max(ys) - min(ys)
        z_range = max(zs) - min(zs)
        
        if volume is None:
            volume = x_range * y_range * z_range
        
        if volume <= 0:
            return {'density': 0}
        
        density = len(self.stations) / volume * 1000  # per 1000 m³
        
        return {
            'stations_per_1000m3': round(density, 2),
            'survey_volume_m3': round(volume),
            'x_extent': round(x_range, 1),
            'y_extent': round(y_range, 1),
            'z_extent': round(z_range, 1)
        }

    def identify_hubs(self, threshold: float = 3) -> List[Dict[str, Union[str, int, float]]]:
        """
        Identify hub stations with high connectivity.
        
        Args:
            threshold: Minimum degree for hub identification
            
        Returns:
            List of hub stations
        """
        hubs = []
        
        for station in self.stations.values():
            degree = len(station.connections)
            if degree >= threshold:
                hubs.append({
                    'name': station.name,
                    'connections': degree,
                    'x': round(station.x, 1),
                    'y': round(station.y, 1),
                    'z': round(station.z, 1),
                    'quality': station.quality
                })
        
        return sorted(hubs, key=lambda x: x['connections'], reverse=True)

    def calculate_quality_metrics(self) -> Dict[str, float]:
        """
        Calculate overall survey quality metrics.
        
        Returns:
            Dictionary with quality metrics
        """
        if not self.stations:
            return {'overall_quality': 0}
        
        # Average station quality
        qualities = [s.quality for s in self.stations.values()]
        mean_quality = np.mean(qualities)
        
        # Network redundancy
        if len(self.graph) > 1:
            n_nodes = self.graph.number_of_nodes()
            n_edges = self.graph.number_of_edges()
            # Minimum edges for tree = n_nodes - 1
            redundancy = (n_edges - (n_nodes - 1)) / (n_nodes - 1) if n_nodes > 1 else 0
        else:
            redundancy = 0
        
        # Closure error quality (inverse of mean error)
        closure_errors = self.calculate_closure_errors()
        if closure_errors:
            mean_error = np.mean([e['error_3d_m'] for e in closure_errors])
            closure_quality = max(0, 1 - mean_error / 10)  # 10m = 0 quality
        else:
            closure_quality = 1.0
        
        # Overall quality (weighted average)
        overall = (mean_quality * 0.4 + 
                  min(1, redundancy) * 0.3 + 
                  closure_quality * 0.3)
        
        return {
            'mean_station_quality': round(mean_quality, 2),
            'network_redundancy': round(redundancy, 2),
            'closure_quality': round(closure_quality, 2),
            'overall_quality': round(overall, 2)
        }

    def get_station_summary(self, station_name: str) -> Dict[str, Union[str, float, List]]:
        """
        Get detailed summary for a specific station.
        
        Args:
            station_name: Name of station
            
        Returns:
            Dictionary with station details
        """
        if station_name not in self.stations:
            return {'error': f'Station {station_name} not found'}
        
        station = self.stations[station_name]
        
        # Find distances to other stations
        distances = []
        for other_name, other in self.stations.items():
            if other_name != station_name:
                dist = np.sqrt((station.x - other.x)**2 + 
                             (station.y - other.y)**2 + 
                             (station.z - other.z)**2)
                distances.append((other_name, dist))
        
        distances.sort(key=lambda x: x[1])
        
        return {
            'name': station.name,
            'coordinates': (round(station.x, 2), round(station.y, 2), round(station.z, 2)),
            'connections': station.connections,
            'n_connections': len(station.connections),
            'quality': station.quality,
            'nearest_station': distances[0] if distances else None,
            'farthest_station': distances[-1] if distances else None
        }

    def summary(self) -> Dict[str, Union[int, float, str, Dict]]:
        """Get complete summary of station analysis."""
        if not self.stations:
            return {'status': 'No stations'}
        
        network = self.analyze_network()
        quality = self.calculate_quality_metrics()
        hubs = self.identify_hubs()
        isolated = self.find_isolated_stations()
        
        return {
            'n_stations': len(self.stations),
            'network_analysis': network,
            'quality_metrics': quality,
            'hub_stations': hubs[:5],  # Top 5 hubs
            'isolated_stations': isolated[:5],  # Top 5 isolated
            'major_closure_errors': self.calculate_closure_errors()[:3]  # Top 3 errors
        }


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Station Analysis")
    print("="*60)
    
    sa = StationAnalysis()
    
    # Add stations
    sa.add_station("A1", 0, 0, 0, quality=0.95)
    sa.add_station("A2", 100, 0, -10, quality=0.90)
    sa.add_station("A3", 150, 50, -25, quality=0.85)
    sa.add_station("B1", 200, 100, -50, quality=0.80)
    
    # Add connections
    sa.add_connection("A1", "A2", 100)
    sa.add_connection("A2", "A3", 70)
    sa.add_connection("A3", "B1", 80)
    sa.add_connection("A1", "B1", 250)  # Long connection
    
    # Analyze network
    network = sa.analyze_network()
    print(f"\nNetwork Analysis:")
    print(f"  - Stations: {network['n_stations']}")
    print(f"  - Shots: {network['n_shots']}")
    print(f"  - Mean Degree: {network['mean_degree']}")
    print(f"  - Most Central: {network['most_central_station']}")
    
    # Quality metrics
    quality = sa.calculate_quality_metrics()
    print(f"\nQuality Metrics:")
    print(f"  - Overall Quality: {quality['overall_quality']:.2f}")
    
    # Hub stations
    hubs = sa.identify_hubs()
    print(f"\nHub Stations:")
    for hub in hubs:
        print(f"  - {hub['name']}: {hub['connections']} connections")
    
    print("\n" + "="*60)
