"""
Passage 3D modeling and reconstruction from survey data.
Creates 3D models of cave passages for visualization and analysis.

Author: Samir Baladi
"""

import numpy as np
from typing import Dict, List, Tuple, Optional, Union
from dataclasses import dataclass
import logging
from scipy.interpolate import interp1d, UnivariateSpline

logger = logging.getLogger(__name__)

@dataclass
class PassageCrossSection:
    """Represents a passage cross-section."""
    station: str
    x: float
    y: float
    z: float
    width: float
    height: float
    area: float
    shape: str  # 'circle', 'ellipse', 'keyhole', 'canyon'


class PassageModeling:
    """
    Creates 3D models of cave passages from survey data.
    """
    
    SHAPE_TYPES = ['circle', 'ellipse', 'keyhole', 'canyon', 'rift', 'tube']
    
    def __init__(self):
        """Initialize passage modeling."""
        self.sections = []
        self.centerline = []
        self.surface_points = []
        logger.info("Initialized PassageModeling")

    def add_cross_section(self,
                         station: str,
                         x: float,
                         y: float,
                         z: float,
                         left: float,
                         right: float,
                         up: float,
                         down: float) -> None:
        """
        Add a passage cross-section from LRUD measurements.
        
        Args:
            station: Station name
            x, y, z: Station coordinates
            left, right: Left and right wall distances
            up, down: Ceiling and floor distances
        """
        width = left + right
        height = up + down
        
        # Estimate cross-sectional area (simplified ellipse)
        area = np.pi * (width/2) * (height/2)
        
        # Determine shape based on aspect ratio
        aspect = height / width if width > 0 else 1
        
        if aspect > 2:
            shape = 'canyon'
        elif aspect < 0.5:
            shape = 'tube'
        elif 0.8 < aspect < 1.2:
            shape = 'circle'
        else:
            shape = 'ellipse'
        
        section = PassageCrossSection(
            station=station,
            x=x, y=y, z=z,
            width=width,
            height=height,
            area=area,
            shape=shape
        )
        
        self.sections.append(section)
        self.centerline.append((x, y, z))
        
        # Generate surface points for this section
        self._generate_section_points(section, left, right, up, down)
        
        logger.debug(f"Added section at {station}: {width:.1f}m x {height:.1f}m")

    def _generate_section_points(self,
                                section: PassageCrossSection,
                                left: float,
                                right: float,
                                up: float,
                                down: float,
                                n_points: int = 16) -> None:
        """
        Generate surface points around the cross-section.
        
        Args:
            section: Cross-section data
            left, right, up, down: Wall distances
            n_points: Number of points to generate
        """
        # Generate points around the perimeter
        angles = np.linspace(0, 2*np.pi, n_points, endpoint=False)
        
        # Simple elliptical shape
        for angle in angles:
            # Radial distance varies with angle
            if angle < np.pi/2:  # Top-right quadrant
                r_x = right * np.cos(angle)
                r_y = up * np.sin(angle)
            elif angle < np.pi:  # Top-left quadrant
                r_x = -left * np.cos(np.pi - angle)
                r_y = up * np.sin(np.pi - angle)
            elif angle < 3*np.pi/2:  # Bottom-left quadrant
                r_x = -left * np.cos(angle - np.pi)
                r_y = -down * np.sin(angle - np.pi)
            else:  # Bottom-right quadrant
                r_x = right * np.cos(2*np.pi - angle)
                r_y = -down * np.sin(2*np.pi - angle)
            
            # Transform to world coordinates
            # Assuming passage oriented along centerline direction
            # This is simplified - in practice need orientation from survey
            point = (
                section.x + r_x,
                section.y + r_y,
                section.z
            )
            
            self.surface_points.append(point)

    def interpolate_centerline(self, n_points: int = 100) -> np.ndarray:
        """
        Interpolate smooth centerline through stations.
        
        Args:
            n_points: Number of interpolation points
            
        Returns:
            Array of interpolated points
        """
        if len(self.centerline) < 2:
            return np.array(self.centerline)
        
        points = np.array(self.centerline)
        t = np.linspace(0, 1, len(points))
        
        # Cubic spline interpolation
        from scipy.interpolate import CubicSpline
        
        cs_x = CubicSpline(t, points[:, 0])
        cs_y = CubicSpline(t, points[:, 1])
        cs_z = CubicSpline(t, points[:, 2])
        
        t_new = np.linspace(0, 1, n_points)
        
        interpolated = np.column_stack([
            cs_x(t_new),
            cs_y(t_new),
            cs_z(t_new)
        ])
        
        return interpolated

    def create_mesh(self, resolution: float = 1.0) -> Dict[str, np.ndarray]:
        """
        Create a 3D mesh of the passage.
        
        Args:
            resolution: Mesh resolution (m)
            
        Returns:
            Dictionary with vertices and faces
        """
        if len(self.sections) < 2:
            return {'vertices': np.array([]), 'faces': np.array([])}
        
        vertices = []
        faces = []
        
        # Sort sections by position along passage
        sections_sorted = sorted(self.sections, key=lambda s: s.z)
        
        # Generate vertices for each section
        for i, section in enumerate(sections_sorted):
            # Generate points around section perimeter
            section_vertices = self._generate_section_vertices(section, resolution)
            vertices.extend(section_vertices)
            
            # Create faces between this section and next
            if i < len(sections_sorted) - 1:
                next_section = sections_sorted[i + 1]
                next_vertices = self._generate_section_vertices(next_section, resolution)
                
                # Connect corresponding vertices
                n_verts = len(section_vertices)
                for j in range(n_verts):
                    v1 = i * n_verts + j
                    v2 = i * n_verts + (j + 1) % n_verts
                    v3 = (i + 1) * n_verts + j
                    v4 = (i + 1) * n_verts + (j + 1) % n_verts
                    
                    # Two triangles per quad
                    faces.append([v1, v2, v3])
                    faces.append([v2, v4, v3])
        
        return {
            'vertices': np.array(vertices),
            'faces': np.array(faces)
        }

    def _generate_section_vertices(self,
                                  section: PassageCrossSection,
                                  resolution: float) -> List[Tuple[float, float, float]]:
        """Generate vertices for a single cross-section."""
        vertices = []
        
        # Number of points based on perimeter
        perimeter = 2 * np.pi * np.sqrt((section.width/2)**2 + (section.height/2)**2) / 2
        n_points = max(8, int(perimeter / resolution))
        
        for i in range(n_points):
            angle = 2 * np.pi * i / n_points
            
            # Elliptical shape
            rx = (section.width / 2) * np.cos(angle)
            ry = (section.height / 2) * np.sin(angle)
            
            # Transform to 3D (simplified - assumes passage aligned with z-axis)
            x = section.x + rx
            y = section.y + ry
            z = section.z
            
            vertices.append((x, y, z))
        
        return vertices

    def calculate_volume(self) -> float:
        """
        Calculate passage volume.
        
        Returns:
            Total volume (m³)
        """
        if len(self.sections) < 2:
            return 0
        
        volume = 0
        sections_sorted = sorted(self.sections, key=lambda s: s.z)
        
        for i in range(len(sections_sorted) - 1):
            s1 = sections_sorted[i]
            s2 = sections_sorted[i + 1]
            
            # Distance between sections
            dist = np.sqrt((s2.x - s1.x)**2 + (s2.y - s1.y)**2 + (s2.z - s1.z)**2)
            
            # Average area
            avg_area = (s1.area + s2.area) / 2
            
            # Volume of this segment
            volume += avg_area * dist
        
        return volume

    def identify_passage_type(self,
                             section: PassageCrossSection) -> Dict[str, Union[str, float]]:
        """
        Identify passage type based on morphology.
        
        Args:
            section: Cross-section data
            
        Returns:
            Dictionary with passage type classification
        """
        aspect = section.height / section.width if section.width > 0 else 1
        
        if section.shape == 'canyon' or aspect > 2:
            ptype = "Vadose Canyon"
            description = "Vertical downcutting passage"
            formation = "Flowing water incision"
        elif section.shape == 'tube' or aspect < 0.5:
            ptype = "Phreatic Tube"
            description = "Circular water-filled conduit"
            formation = "Dissolution under pressure"
        elif section.shape == 'circle' or 0.8 < aspect < 1.2:
            ptype = "Phreatic Conduit"
            description = "Elliptical water passage"
            formation = "Dissolutional enlargement"
        else:
            ptype = "Mixed Morphology"
            description = "Multiple speleogenetic phases"
            formation = "Complex history"
        
        return {
            'type': ptype,
            'description': description,
            'formation': formation,
            'aspect_ratio': round(aspect, 2),
            'width': round(section.width, 1),
            'height': round(section.height, 1)
        }

    def calculate_hydraulic_radius(self, section: PassageCrossSection) -> float:
        """
        Calculate hydraulic radius for flow analysis.
        
        Args:
            section: Cross-section data
            
        Returns:
            Hydraulic radius (m)
        """
        # Area / wetted perimeter (simplified for ellipse)
        area = section.area
        
        # Approximate wetted perimeter for ellipse
        a = section.width / 2
        b = section.height / 2
        perimeter = np.pi * (3*(a + b) - np.sqrt((3*a + b)*(a + 3*b)))
        
        if perimeter > 0:
            return area / perimeter
        else:
            return 0

    def export_to_stl(self, filename: str) -> None:
        """
        Export 3D model to STL file.
        
        Args:
            filename: Output filename
        """
        mesh = self.create_mesh()
        
        if len(mesh['vertices']) == 0:
            logger.warning("No mesh to export")
            return
        
        with open(filename, 'w') as f:
            f.write("solid CAVORA_Passage\n")
            
            for face in mesh['faces']:
                f.write(f"  facet normal 0 0 0\n")
                f.write("    outer loop\n")
                for vi in face:
                    v = mesh['vertices'][vi]
                    f.write(f"      vertex {v[0]:.3f} {v[1]:.3f} {v[2]:.3f}\n")
                f.write("    endloop\n")
                f.write("  endfacet\n")
            
            f.write("endsolid CAVORA_Passage\n")
        
        logger.info(f"Exported STL to {filename}")

    def summary(self) -> Dict[str, Union[int, float, str]]:
        """Get summary of passage model."""
        if not self.sections:
            return {'status': 'No sections'}
        
        widths = [s.width for s in self.sections]
        heights = [s.height for s in self.sections]
        areas = [s.area for s in self.sections]
        
        # Calculate passage length
        if len(self.centerline) > 1:
            length = 0
            for i in range(len(self.centerline) - 1):
                p1 = self.centerline[i]
                p2 = self.centerline[i + 1]
                length += np.sqrt((p2[0]-p1[0])**2 + (p2[1]-p1[1])**2 + (p2[2]-p1[2])**2)
        else:
            length = 0
        
        volume = self.calculate_volume()
        
        return {
            'n_sections': len(self.sections),
            'length_m': round(length, 1),
            'volume_m3': round(volume),
            'mean_width': round(np.mean(widths), 2),
            'mean_height': round(np.mean(heights), 2),
            'mean_area': round(np.mean(areas), 1),
            'max_width': round(max(widths), 2),
            'max_height': round(max(heights), 2),
            'shape_distribution': self._get_shape_distribution()
        }

    def _get_shape_distribution(self) -> Dict[str, int]:
        """Get distribution of passage shapes."""
        shapes = {}
        for section in self.sections:
            shapes[section.shape] = shapes.get(section.shape, 0) + 1
        return shapes


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Passage 3D Modeling")
    print("="*60)
    
    pm = PassageModeling()
    
    # Add some cross-sections
    pm.add_cross_section("A1", 0, 0, 0, 2.0, 3.0, 4.0, 2.0)
    pm.add_cross_section("A2", 10, 5, -5, 2.5, 3.5, 4.5, 2.5)
    pm.add_cross_section("A3", 20, 10, -10, 3.0, 4.0, 5.0, 3.0)
    
    # Get summary
    summary = pm.summary()
    print(f"\nPassage Summary:")
    print(f"  - Length: {summary['length_m']}m")
    print(f"  - Volume: {summary['volume_m3']}m³")
    print(f"  - Mean Width: {summary['mean_width']}m")
    print(f"  - Mean Height: {summary['mean_height']}m")
    
    # Identify passage type
    ptype = pm.identify_passage_type(pm.sections[0])
    print(f"\nPassage Type:")
    print(f"  - Type: {ptype['type']}")
    print(f"  - Description: {ptype['description']}")
    
    print("\n" + "="*60)
