"""
Compass/Therion data parser for cave survey files.
Imports and processes survey data from various formats.

Author: Samir Baladi
"""

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

logger = logging.getLogger(__name__)

@dataclass
class SurveyShot:
    """Represents a single survey shot."""
    from_station: str
    to_station: str
    length: float  # meters
    azimuth: float  # degrees
    inclination: float  # degrees
    left: float  # left wall distance (m)
    right: float  # right wall distance (m)
    up: float  # ceiling distance (m)
    down: float  # floor distance (m)
    comment: str = ""


@dataclass
class SurveyStation:
    """Represents a survey station with coordinates."""
    name: str
    x: float  # meters
    y: float  # meters
    z: float  # meters
    shots: List[SurveyShot]


class CompassParser:
    """
    Parses Compass/THERION survey data files.
    Supports .dat (Compass) and .th (Therion) formats.
    """
    
    def __init__(self):
        """Initialize the parser."""
        self.stations = {}
        self.shots = []
        self.survey_name = ""
        self.units = "meters"
        logger.info("Initialized CompassParser")

    def parse_file(self, filename: str) -> Dict[str, Union[str, List, Dict]]:
        """
        Parse a survey file based on extension.
        
        Args:
            filename: Path to survey file
            
        Returns:
            Dictionary with parsed data
        """
        filepath = Path(filename)
        
        if not filepath.exists():
            raise FileNotFoundError(f"File not found: {filename}")
        
        if filepath.suffix.lower() == '.dat':
            return self._parse_compass(filepath)
        elif filepath.suffix.lower() == '.th':
            return self._parse_therion(filepath)
        else:
            raise ValueError(f"Unsupported file format: {filepath.suffix}")

    def _parse_compass(self, filepath: Path) -> Dict[str, Union[str, List, Dict]]:
        """
        Parse Compass .dat format.
        
        Format: FROM TO LENGTH AZIMUTH INCLINATION LEFT RIGHT UP DOWN COMMENT
        """
        logger.info(f"Parsing Compass file: {filepath}")
        
        with open(filepath, 'r') as f:
            lines = f.readlines()
        
        # First line might be survey name
        if lines and lines[0].startswith('#'):
            self.survey_name = lines[0][1:].strip()
            lines = lines[1:]
        
        for line_num, line in enumerate(lines, 1):
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            
            # Parse shot data
            parts = line.split()
            if len(parts) < 7:
                logger.warning(f"Skipping line {line_num}: insufficient data")
                continue
            
            try:
                shot = SurveyShot(
                    from_station=parts[0],
                    to_station=parts[1],
                    length=float(parts[2]),
                    azimuth=float(parts[3]),
                    inclination=float(parts[4]),
                    left=float(parts[5]) if len(parts) > 5 else 0,
                    right=float(parts[6]) if len(parts) > 6 else 0,
                    up=float(parts[7]) if len(parts) > 7 else 0,
                    down=float(parts[8]) if len(parts) > 8 else 0,
                    comment=' '.join(parts[9:]) if len(parts) > 9 else ""
                )
                
                self.shots.append(shot)
                
            except (ValueError, IndexError) as e:
                logger.error(f"Error parsing line {line_num}: {e}")
        
        # Calculate station coordinates
        self._calculate_coordinates()
        
        return self.get_survey_data()

    def _parse_therion(self, filepath: Path) -> Dict[str, Union[str, List, Dict]]:
        """
        Parse Therion .th format.
        
        Basic format: centerline ... endcenterline
        """
        logger.info(f"Parsing Therion file: {filepath}")
        
        with open(filepath, 'r') as f:
            content = f.read()
        
        # Extract survey name
        name_match = re.search(r'survey\s+(\w+)', content)
        if name_match:
            self.survey_name = name_match.group(1)
        
        # Extract centerline data
        centerline_pattern = r'centerline.*?\n(.*?)endcenterline'
        centerline_matches = re.findall(centerline_pattern, content, re.DOTALL | re.IGNORECASE)
        
        for centerline in centerline_matches:
            lines = centerline.strip().split('\n')
            
            for line in lines:
                line = line.strip()
                if not line or line.startswith('#'):
                    continue
                
                # Parse shot: from to length azimuth inclination
                parts = line.split()
                if len(parts) < 5:
                    continue
                
                try:
                    # Check if it's a shot line
                    if parts[0] in ['from', 'to', 'data']:
                        continue
                    
                    shot = SurveyShot(
                        from_station=parts[0],
                        to_station=parts[1],
                        length=float(parts[2]),
                        azimuth=float(parts[3]),
                        inclination=float(parts[4]),
                        left=0, right=0, up=0, down=0  # Therion stores dimensions separately
                    )
                    
                    self.shots.append(shot)
                    
                except (ValueError, IndexError) as e:
                    logger.warning(f"Error parsing Therion line: {line}")
        
        # Calculate coordinates
        self._calculate_coordinates()
        
        return self.get_survey_data()

    def _calculate_coordinates(self):
        """Calculate station coordinates from shots."""
        # Start with first station at origin
        if not self.shots:
            return
        
        # Use first shot's from_station as origin
        origin = self.shots[0].from_station
        self.stations[origin] = SurveyStation(
            name=origin,
            x=0, y=0, z=0,
            shots=[]
        )
        
        # Process all shots
        shots_processed = set()
        while len(shots_processed) < len(self.shots):
            for shot in self.shots:
                if id(shot) in shots_processed:
                    continue
                
                # Check if we know the from_station
                if shot.from_station in self.stations:
                    from_station = self.stations[shot.from_station]
                    
                    # Calculate new coordinates
                    rad_az = np.radians(shot.azimuth)
                    rad_inc = np.radians(shot.inclination)
                    
                    # Convert polar to Cartesian
                    dx = shot.length * np.cos(rad_inc) * np.sin(rad_az)
                    dy = shot.length * np.cos(rad_inc) * np.cos(rad_az)
                    dz = shot.length * np.sin(rad_inc)
                    
                    new_x = from_station.x + dx
                    new_y = from_station.y + dy
                    new_z = from_station.z + dz
                    
                    # Create or update station
                    if shot.to_station in self.stations:
                        # Check for closure error
                        existing = self.stations[shot.to_station]
                        error = np.sqrt((new_x - existing.x)**2 + 
                                      (new_y - existing.y)**2 + 
                                      (new_z - existing.z)**2)
                        if error > 0.1:  # 10cm tolerance
                            logger.warning(f"Closure error at {shot.to_station}: {error:.3f}m")
                    else:
                        self.stations[shot.to_station] = SurveyStation(
                            name=shot.to_station,
                            x=new_x, y=new_y, z=new_z,
                            shots=[]
                        )
                    
                    # Add shot to both stations
                    self.stations[shot.from_station].shots.append(shot)
                    if shot.to_station in self.stations:
                        self.stations[shot.to_station].shots.append(shot)
                    
                    shots_processed.add(id(shot))
                    
                elif shot.to_station in self.stations:
                    # Reverse shot
                    to_station = self.stations[shot.to_station]
                    
                    rad_az = np.radians((shot.azimuth + 180) % 360)
                    rad_inc = np.radians(-shot.inclination)
                    
                    dx = shot.length * np.cos(rad_inc) * np.sin(rad_az)
                    dy = shot.length * np.cos(rad_inc) * np.cos(rad_az)
                    dz = shot.length * np.sin(rad_inc)
                    
                    new_x = to_station.x + dx
                    new_y = to_station.y + dy
                    new_z = to_station.z + dz
                    
                    if shot.from_station in self.stations:
                        existing = self.stations[shot.from_station]
                        error = np.sqrt((new_x - existing.x)**2 + 
                                      (new_y - existing.y)**2 + 
                                      (new_z - existing.z)**2)
                        if error > 0.1:
                            logger.warning(f"Closure error at {shot.from_station}: {error:.3f}m")
                    else:
                        self.stations[shot.from_station] = SurveyStation(
                            name=shot.from_station,
                            x=new_x, y=new_y, z=new_z,
                            shots=[]
                        )
                    
                    shots_processed.add(id(shot))

    def get_survey_data(self) -> Dict[str, Union[str, List, Dict]]:
        """Get all survey data."""
        stations_list = []
        for name, station in self.stations.items():
            stations_list.append({
                'name': name,
                'x': round(station.x, 2),
                'y': round(station.y, 2),
                'z': round(station.z, 2),
                'n_shots': len(station.shots)
            })
        
        shots_list = []
        for shot in self.shots:
            shots_list.append({
                'from': shot.from_station,
                'to': shot.to_station,
                'length': shot.length,
                'azimuth': shot.azimuth,
                'inclination': shot.inclination,
                'lrdu': [shot.left, shot.right, shot.up, shot.down]
            })
        
        return {
            'survey_name': self.survey_name,
            'n_stations': len(self.stations),
            'n_shots': len(self.shots),
            'total_length': sum(s.length for s in self.shots),
            'stations': stations_list,
            'shots': shots_list,
            'units': self.units
        }

    def export_to_therion(self, filename: str) -> None:
        """
        Export survey data to Therion format.
        
        Args:
            filename: Output filename
        """
        with open(filename, 'w') as f:
            f.write(f"survey {self.survey_name or 'cave'} -export\n")
            f.write("centerline\n")
            
            for shot in self.shots:
                f.write(f"  {shot.from_station} {shot.to_station} "
                       f"{shot.length:.2f} {shot.azimuth:.1f} {shot.inclination:.1f}")
                if shot.comment:
                    f.write(f"  ; {shot.comment}")
                f.write("\n")
            
            f.write("endcenterline\n")
            f.write("endsurvey\n")
        
        logger.info(f"Exported {len(self.shots)} shots to {filename}")

    def export_to_compass(self, filename: str) -> None:
        """
        Export survey data to Compass .dat format.
        
        Args:
            filename: Output filename
        """
        with open(filename, 'w') as f:
            if self.survey_name:
                f.write(f"# {self.survey_name}\n")
            
            for shot in self.shots:
                f.write(f"{shot.from_station:<6} {shot.to_station:<6} "
                       f"{shot.length:6.2f} {shot.azimuth:6.1f} {shot.inclination:5.1f} "
                       f"{shot.left:4.2f} {shot.right:4.2f} {shot.up:4.2f} {shot.down:4.2f}")
                if shot.comment:
                    f.write(f"  ; {shot.comment}")
                f.write("\n")
        
        logger.info(f"Exported {len(self.shots)} shots to {filename}")

    def summary(self) -> Dict[str, Union[str, float, int]]:
        """Get summary of survey data."""
        if not self.shots:
            return {'status': 'No survey data'}
        
        lengths = [s.length for s in self.shots]
        
        return {
            'survey_name': self.survey_name or 'Unnamed',
            'n_stations': len(self.stations),
            'n_shots': len(self.shots),
            'total_length_m': round(sum(lengths), 1),
            'mean_shot_length': round(np.mean(lengths), 2),
            'max_shot_length': round(max(lengths), 2),
            'min_shot_length': round(min(lengths), 2),
            'x_range': self._get_coordinate_range('x'),
            'y_range': self._get_coordinate_range('y'),
            'z_range': self._get_coordinate_range('z')
        }

    def _get_coordinate_range(self, coord: str) -> str:
        """Get range of coordinates."""
        if not self.stations:
            return "0-0"
        
        values = [getattr(s, coord) for s in self.stations.values()]
        return f"{min(values):.1f} - {max(values):.1f}"


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Compass/Therion Parser")
    print("="*60)
    
    # Create parser
    parser = CompassParser()
    
    # Add some sample shots manually
    parser.shots.append(SurveyShot("A1", "A2", 15.5, 45.0, 5.0, 2.0, 3.0, 4.0, 2.0))
    parser.shots.append(SurveyShot("A2", "A3", 22.3, 90.0, -2.0, 3.0, 4.0, 5.0, 3.0))
    
    # Calculate coordinates
    parser._calculate_coordinates()
    
    # Get summary
    summary = parser.summary()
    print(f"\nSurvey Summary:")
    print(f"  - Name: {summary['survey_name']}")
    print(f"  - Stations: {summary['n_stations']}")
    print(f"  - Shots: {summary['n_shots']}")
    print(f"  - Total Length: {summary['total_length_m']}m")
    print(f"  - Depth Range: {summary['z_range']}m")
    
    print("\n" + "="*60)
