"""
Brace surface generation from spine landmarks.

Two modes:
- Version A: Generic/average body shape (parametric torso)
- Version B: Uses actual 3D body scan mesh
"""
import numpy as np
from typing import Tuple, Optional, List
from pathlib import Path

from .data_models import Spine2D, BraceConfig
from .spine_analysis import compute_spine_curve, find_apex_vertebrae

try:
    import trimesh
    HAS_TRIMESH = True
except ImportError:
    HAS_TRIMESH = False


class BraceGenerator:
    """
    Generates 3D brace shell from spine landmarks.
    """
    
    def __init__(self, config: Optional[BraceConfig] = None):
        """
        Initialize brace generator.
        
        Args:
            config: Brace configuration parameters
        """
        if not HAS_TRIMESH:
            raise ImportError("trimesh is required for brace generation. Install with: pip install trimesh")
        
        self.config = config or BraceConfig()
    
    def generate(self, spine: Spine2D) -> 'trimesh.Trimesh':
        """
        Generate brace mesh from spine landmarks.
        
        Args:
            spine: Spine2D object with detected vertebrae
            
        Returns:
            trimesh.Trimesh object representing the brace shell
        """
        if self.config.use_body_scan and self.config.body_scan_path:
            return self._generate_from_body_scan(spine)
        else:
            return self._generate_from_average_body(spine)
    
    def _torso_profile(self, z01: float) -> Tuple[float, float]:
        """
        Get torso cross-section radii at a given height.
        
        Args:
            z01: Normalized height (0=top, 1=bottom)
            
        Returns:
            (a_mm, b_mm): Radii in left-right and front-back directions
        """
        # Torso shape varies with height
        # Wider at chest (z~0.3) and hips (z~0.8), narrower at waist (z~0.5)
        
        # Base radii from config
        base_a = self.config.torso_width_mm / 2
        base_b = self.config.torso_depth_mm / 2
        
        # Shape modulation
        # Chest region (z ~ 0.2-0.4): wider
        # Waist region (z ~ 0.5): narrower
        # Hip region (z ~ 0.8-1.0): wider
        
        if z01 < 0.3:
            # Upper chest - moderate width
            mod = 1.0
        elif z01 < 0.5:
            # Transition to waist
            t = (z01 - 0.3) / 0.2
            mod = 1.0 - 0.15 * t  # Decrease by 15%
        elif z01 < 0.7:
            # Waist region - narrowest
            mod = 0.85
        else:
            # Hips - widen again
            t = (z01 - 0.7) / 0.3
            mod = 0.85 + 0.2 * t  # Increase by 20%
        
        return base_a * mod, base_b * mod
    
    def _generate_from_average_body(self, spine: Spine2D) -> 'trimesh.Trimesh':
        """
        Generate brace using parametric average body shape.
        
        The brace follows the spine curve laterally and applies
        pressure zones at curve apexes.
        """
        cfg = self.config
        
        # 1) Compute spine curve
        try:
            C_px, T_px, N_px, curvature = compute_spine_curve(spine, smooth=5.0, n_samples=cfg.n_vertical_slices)
        except ValueError as e:
            raise ValueError(f"Cannot generate brace: {e}")
        
        # 2) Convert to mm
        if spine.pixel_spacing_mm is not None:
            sx, sy = spine.pixel_spacing_mm
        elif cfg.pixel_spacing_mm is not None:
            sx, sy = cfg.pixel_spacing_mm
        else:
            sx = sy = 0.25  # Default assumption
        
        C_mm = np.zeros_like(C_px)
        C_mm[:, 0] = C_px[:, 0] * sx
        C_mm[:, 1] = C_px[:, 1] * sy
        
        # 3) Determine brace vertical extent
        y_mm = C_mm[:, 1]
        y_min, y_max = y_mm.min(), y_mm.max()
        spine_height = y_max - y_min
        
        # Brace height (might extend beyond detected vertebrae)
        brace_height = min(cfg.brace_height_mm, spine_height * 1.1)
        
        # 4) Normalize curvature for pressure zones
        curv_norm = (curvature - curvature.min()) / (curvature.max() - curvature.min() + 1e-8)
        
        # 5) Build vertices
        n_z = cfg.n_vertical_slices
        n_theta = cfg.n_circumference_points
        
        # Opening angle (front of brace might be open)
        opening_half = np.radians(cfg.front_opening_deg / 2)
        
        vertices = []
        
        for i in range(n_z):
            z01 = i / (n_z - 1)  # 0 to 1
            
            # Z coordinate (vertical position in 3D)
            z_mm = y_min + z01 * spine_height
            
            # Get torso profile at this height
            a_mm, b_mm = self._torso_profile(z01)
            
            # Lateral offset from spine curve
            x_offset = C_mm[i, 0] - (C_mm[0, 0] + C_mm[-1, 0]) / 2  # Deviation from midline
            
            # Pressure modulation based on curvature
            pressure = cfg.pressure_strength_mm * curv_norm[i]
            
            for j in range(n_theta):
                theta = 2 * np.pi * (j / n_theta)
                
                # Skip vertices in the opening region (front = theta around 0)
                # Actually, we'll still create them but can mark them for later removal
                
                # Base ellipse point
                x = a_mm * np.cos(theta)
                y = b_mm * np.sin(theta)
                
                # Apply lateral offset (brace follows spine curve)
                x += x_offset
                
                # Apply pressure zones
                # Pressure on sides (theta near π/2 or 3π/2 = sides)
                # The side that's convex gets pushed in
                side_factor = abs(np.cos(theta))  # Max at sides (theta=0 or π)
                
                # Determine which side based on spine deviation
                if x_offset > 0:
                    # Spine deviated right, push on right side
                    if np.cos(theta) > 0:  # Right side
                        x -= pressure * side_factor
                else:
                    # Spine deviated left, push on left side
                    if np.cos(theta) < 0:  # Left side
                        x -= pressure * side_factor * np.sign(np.cos(theta))
                
                # Vertex position: x=left/right, y=front/back, z=vertical
                vertices.append([x, y, z_mm])
        
        vertices = np.array(vertices, dtype=np.float32)
        
        # 6) Build faces (quad strips between adjacent rings)
        faces = []
        
        def vid(i, j):
            return i * n_theta + (j % n_theta)
        
        for i in range(n_z - 1):
            for j in range(n_theta):
                j2 = (j + 1) % n_theta
                
                # Two triangles per quad
                a = vid(i, j)
                b = vid(i, j2)
                c = vid(i + 1, j2)
                d = vid(i + 1, j)
                
                faces.append([a, b, c])
                faces.append([a, c, d])
        
        faces = np.array(faces, dtype=np.int32)
        
        # 7) Create outer shell mesh
        outer_shell = trimesh.Trimesh(vertices=vertices, faces=faces, process=True)
        
        # 8) Create inner shell (offset inward by wall thickness)
        outer_shell.fix_normals()
        vn = outer_shell.vertex_normals
        inner_vertices = vertices - cfg.wall_thickness_mm * vn
        
        # Inner faces need reversed winding
        inner_faces = faces[:, ::-1]
        
        # 9) Combine into solid shell
        all_vertices = np.vstack([vertices, inner_vertices])
        inner_faces_offset = inner_faces + len(vertices)
        all_faces = np.vstack([faces, inner_faces_offset])
        
        # 10) Add end caps (top and bottom rings)
        # Top cap (connect outer to inner at z=0)
        top_faces = []
        for j in range(n_theta):
            j2 = (j + 1) % n_theta
            outer_j = vid(0, j)
            outer_j2 = vid(0, j2)
            inner_j = outer_j + len(vertices)
            inner_j2 = outer_j2 + len(vertices)
            top_faces.append([outer_j, inner_j, inner_j2])
            top_faces.append([outer_j, inner_j2, outer_j2])
        
        # Bottom cap
        bottom_faces = []
        for j in range(n_theta):
            j2 = (j + 1) % n_theta
            outer_j = vid(n_z - 1, j)
            outer_j2 = vid(n_z - 1, j2)
            inner_j = outer_j + len(vertices)
            inner_j2 = outer_j2 + len(vertices)
            bottom_faces.append([outer_j, outer_j2, inner_j2])
            bottom_faces.append([outer_j, inner_j2, inner_j])
        
        all_faces = np.vstack([all_faces, top_faces, bottom_faces])
        
        # Create final mesh
        brace = trimesh.Trimesh(vertices=all_vertices, faces=all_faces, process=True)
        brace.merge_vertices()
        # Remove degenerate faces
        valid_faces = brace.nondegenerate_faces()
        brace.update_faces(valid_faces)
        brace.fix_normals()
        
        return brace
    
    def _generate_from_body_scan(self, spine: Spine2D) -> 'trimesh.Trimesh':
        """
        Generate brace by offsetting from a 3D body scan mesh.
        
        The body scan provides the actual torso shape, and we:
        1. Offset outward for clearance
        2. Apply pressure zones based on spine curvature
        3. Thicken for wall thickness
        """
        cfg = self.config
        
        if not cfg.body_scan_path or not Path(cfg.body_scan_path).exists():
            raise FileNotFoundError(f"Body scan not found: {cfg.body_scan_path}")
        
        # Load body scan
        body = trimesh.load(cfg.body_scan_path, force='mesh')
        body.remove_unreferenced_vertices()
        body.fix_normals()
        
        # Compute spine curve for pressure mapping
        try:
            C_px, T_px, N_px, curvature = compute_spine_curve(spine, smooth=5.0, n_samples=200)
        except ValueError:
            curvature = np.zeros(200)
        
        # Convert spine coordinates to mm
        if spine.pixel_spacing_mm is not None:
            sx, sy = spine.pixel_spacing_mm
        else:
            sx = sy = 0.25
        
        y_mm = C_px[:, 1] * sy
        y_min, y_max = y_mm.min(), y_mm.max()
        H = y_max - y_min + 1e-6
        
        # Normalize curvature
        curv_norm = (curvature - curvature.min()) / (curvature.max() - curvature.min() + 1e-8)
        
        # 1) Offset body surface outward for clearance (inner brace surface)
        clearance_mm = 6.0  # Gap between body and brace
        vn = body.vertex_normals
        inner_surface = trimesh.Trimesh(
            vertices=body.vertices + clearance_mm * vn,
            faces=body.faces.copy(),
            process=True
        )
        
        # 2) Apply pressure deformation
        # Map each vertex's Z coordinate to spine curvature
        z_coords = inner_surface.vertices[:, 2]  # Assuming Z is vertical
        z_min, z_max = z_coords.min(), z_coords.max()
        z01 = (z_coords - z_min) / (z_max - z_min + 1e-6)
        
        # Sample curvature at each vertex height
        curv_idx = np.clip((z01 * (len(curv_norm) - 1)).astype(int), 0, len(curv_norm) - 1)
        pressure_per_vertex = cfg.pressure_strength_mm * curv_norm[curv_idx]
        
        # Apply pressure on sides (based on X coordinate)
        x_coords = inner_surface.vertices[:, 0]
        x_range = np.abs(x_coords).max() + 1e-6
        side_factor = np.abs(x_coords) / x_range  # 0 at center, 1 at sides
        
        deformation = (pressure_per_vertex * side_factor)[:, np.newaxis] * inner_surface.vertex_normals
        inner_surface.vertices = inner_surface.vertices - deformation
        
        # 3) Create outer surface (offset by wall thickness)
        inner_surface.fix_normals()
        outer_surface = trimesh.Trimesh(
            vertices=inner_surface.vertices + cfg.wall_thickness_mm * inner_surface.vertex_normals,
            faces=inner_surface.faces.copy(),
            process=True
        )
        
        # 4) Combine surfaces
        # For a true solid, we'd need to stitch edges - simplified here
        brace = trimesh.util.concatenate([inner_surface, outer_surface])
        brace.merge_vertices()
        valid_faces = brace.nondegenerate_faces()
        brace.update_faces(valid_faces)
        brace.fix_normals()
        
        return brace
    
    def export_stl(self, mesh: 'trimesh.Trimesh', output_path: str):
        """
        Export mesh to STL file.
        
        Args:
            mesh: trimesh.Trimesh object
            output_path: Path for output STL file
        """
        mesh.export(output_path)
        print(f"Exported brace to {output_path}")
        print(f"  Vertices: {len(mesh.vertices)}")
        print(f"  Faces: {len(mesh.faces)}")
