"""
GLB Brace Generator with Markers

This module generates GLB brace files with embedded markers for editing.
Supports both regular (fitted) and vase-shaped templates.

PRESSURE ZONES EXPLANATION:
===========================

The brace has 4 main pressure/expansion zones that correct spinal curvature:

1. THORACIC PAD (LM_PAD_TH) - PUSH ZONE
   - Location: On the CONVEX side of the thoracic curve (the side that bulges out)
   - Function: Pushes INWARD to correct the thoracic curvature
   - For right thoracic curves: pad is on the RIGHT back
   - Depth: 8-25mm depending on Cobb angle severity

2. THORACIC BAY (LM_BAY_TH) - EXPANSION ZONE  
   - Location: OPPOSITE the thoracic pad (concave side)
   - Function: Creates SPACE for the body to move INTO during correction
   - The ribs/body shift into this space as the pad pushes
   - Clearance: 10-35mm

3. LUMBAR PAD (LM_PAD_LUM) - PUSH ZONE
   - Location: On the CONVEX side of the lumbar curve
   - Function: Pushes INWARD to correct lumbar curvature
   - Usually on the opposite side of thoracic pad (for S-curves)
   - Depth: 6-20mm

4. LUMBAR BAY (LM_BAY_LUM) - EXPANSION ZONE
   - Location: OPPOSITE the lumbar pad
   - Function: Creates SPACE for lumbar correction
   - Clearance: 8-25mm

5. HIP ANCHORS (LM_ANCHOR_HIP_L/R) - STABILITY ZONES
   - Location: Around the hip/pelvis area on both sides
   - Function: Grip the pelvis to prevent brace from riding up
   - Slight inward pressure to anchor the brace

The Rigo classification determines which zones are primary:
- A types (3-curve): Strong thoracic pad, minor lumbar
- B types (4-curve): Both thoracic and lumbar pads are primary
- C types (non-3-non-4): Balanced thoracic, neutral pelvis
- E types (single lumbar/TL): Strong lumbar/TL pad, counter-thoracic

"""

import json
import numpy as np
import trimesh
from pathlib import Path
from typing import Dict, Any, Optional, Tuple, Literal
from dataclasses import dataclass, asdict

# Paths to template directories
BASE_DIR = Path(__file__).parent.parent
BRACES_DIR = BASE_DIR / "braces"
REGULAR_TEMPLATES_DIR = BRACES_DIR / "brace_templates"
VASE_TEMPLATES_DIR = BRACES_DIR / "vase_brace_templates"

# Template types
TemplateType = Literal["regular", "vase"]

@dataclass
class MarkerPositions:
    """Marker positions for a brace template."""
    LM_PELVIS_CENTER: Tuple[float, float, float]
    LM_TOP_CENTER: Tuple[float, float, float]
    LM_PAD_TH: Tuple[float, float, float]
    LM_BAY_TH: Tuple[float, float, float]
    LM_PAD_LUM: Tuple[float, float, float]
    LM_BAY_LUM: Tuple[float, float, float]
    LM_ANCHOR_HIP_L: Tuple[float, float, float]
    LM_ANCHOR_HIP_R: Tuple[float, float, float]

@dataclass
class PressureZone:
    """Describes a pressure or expansion zone on the brace."""
    name: str
    marker_name: str
    position: Tuple[float, float, float]
    zone_type: Literal["pad", "bay", "anchor"]
    function: str
    direction: Literal["inward", "outward", "grip"]
    depth_mm: float = 0.0
    radius_mm: Tuple[float, float, float] = (50.0, 80.0, 40.0)

@dataclass
class BraceGenerationResult:
    """Result of brace generation with markers."""
    glb_path: str
    stl_path: str
    json_path: str
    template_type: str
    rigo_type: str
    markers: Dict[str, Tuple[float, float, float]]
    basis: Dict[str, Any]
    pressure_zones: list
    mesh_stats: Dict[str, int]
    transform_applied: Optional[Dict[str, Any]] = None


def get_template_paths(rigo_type: str, template_type: TemplateType) -> Tuple[Path, Path]:
    """
    Get paths to GLB template and markers JSON.
    
    Args:
        rigo_type: Rigo classification (A1, A2, A3, B1, B2, C1, C2, E1, E2)
        template_type: "regular" or "vase"
        
    Returns:
        Tuple of (glb_path, markers_json_path)
    """
    if template_type == "regular":
        glb_path = REGULAR_TEMPLATES_DIR / f"{rigo_type}_marked_v3.glb"
        json_path = REGULAR_TEMPLATES_DIR / f"{rigo_type}_marked_v3.markers.json"
    else:  # vase
        glb_path = VASE_TEMPLATES_DIR / "glb" / f"{rigo_type}_vase_marked.glb"
        json_path = VASE_TEMPLATES_DIR / "markers_json" / f"{rigo_type}_vase_marked.markers.json"
    
    return glb_path, json_path


def load_template_markers(rigo_type: str, template_type: TemplateType) -> Dict[str, Any]:
    """Load markers from JSON file for a template."""
    _, json_path = get_template_paths(rigo_type, template_type)
    
    if not json_path.exists():
        raise FileNotFoundError(f"Markers JSON not found: {json_path}")
    
    with open(json_path, "r") as f:
        return json.load(f)


def load_glb_template(rigo_type: str, template_type: TemplateType) -> trimesh.Trimesh:
    """Load GLB template as trimesh."""
    glb_path, _ = get_template_paths(rigo_type, template_type)
    
    if not glb_path.exists():
        raise FileNotFoundError(f"GLB template not found: {glb_path}")
    
    scene = trimesh.load(str(glb_path))
    
    # If it's a scene, concatenate all meshes
    if isinstance(scene, trimesh.Scene):
        meshes = [g for g in scene.geometry.values() if isinstance(g, trimesh.Trimesh)]
        if meshes:
            mesh = trimesh.util.concatenate(meshes)
        else:
            raise ValueError(f"No valid meshes found in GLB: {glb_path}")
    else:
        mesh = scene
    
    return mesh


def calculate_pressure_zones(
    markers: Dict[str, Any],
    rigo_type: str,
    cobb_angles: Dict[str, float]
) -> list:
    """
    Calculate pressure zone parameters based on markers and analysis.
    
    Args:
        markers: Marker positions from template
        rigo_type: Rigo classification
        cobb_angles: Dict with PT, MT, TL Cobb angles
        
    Returns:
        List of PressureZone objects
    """
    marker_pos = markers.get("markers", markers)
    
    # Calculate severity from Cobb angles
    mt_angle = cobb_angles.get("MT", 0)
    tl_angle = cobb_angles.get("TL", 0)
    
    # Severity mapping: Cobb -> depth
    def cobb_to_depth(angle: float, min_depth: float = 6.0, max_depth: float = 22.0) -> float:
        severity = min(max((angle - 10) / 40, 0), 1)  # 0-1 range
        return min_depth + severity * (max_depth - min_depth)
    
    th_depth = cobb_to_depth(mt_angle, 8.0, 22.0)
    lum_depth = cobb_to_depth(tl_angle, 6.0, 18.0)
    
    # Bay clearance is typically 1.2-1.5x pad depth
    th_clearance = th_depth * 1.3 + 5
    lum_clearance = lum_depth * 1.3 + 4
    
    zones = [
        PressureZone(
            name="Thoracic Pad",
            marker_name="LM_PAD_TH",
            position=tuple(marker_pos.get("LM_PAD_TH", [0, 0, 0])),
            zone_type="pad",
            function="Pushes INWARD on thoracic curve convex side to correct curvature",
            direction="inward",
            depth_mm=th_depth,
            radius_mm=(50.0, 90.0, 40.0)
        ),
        PressureZone(
            name="Thoracic Bay",
            marker_name="LM_BAY_TH",
            position=tuple(marker_pos.get("LM_BAY_TH", [0, 0, 0])),
            zone_type="bay",
            function="Creates SPACE on thoracic concave side for body to shift into",
            direction="outward",
            depth_mm=th_clearance,
            radius_mm=(65.0, 110.0, 55.0)
        ),
        PressureZone(
            name="Lumbar Pad",
            marker_name="LM_PAD_LUM",
            position=tuple(marker_pos.get("LM_PAD_LUM", [0, 0, 0])),
            zone_type="pad",
            function="Pushes INWARD on lumbar curve convex side to correct curvature",
            direction="inward",
            depth_mm=lum_depth,
            radius_mm=(55.0, 85.0, 45.0)
        ),
        PressureZone(
            name="Lumbar Bay",
            marker_name="LM_BAY_LUM",
            position=tuple(marker_pos.get("LM_BAY_LUM", [0, 0, 0])),
            zone_type="bay",
            function="Creates SPACE on lumbar concave side for body to shift into",
            direction="outward",
            depth_mm=lum_clearance,
            radius_mm=(70.0, 100.0, 60.0)
        ),
        PressureZone(
            name="Left Hip Anchor",
            marker_name="LM_ANCHOR_HIP_L",
            position=tuple(marker_pos.get("LM_ANCHOR_HIP_L", [0, 0, 0])),
            zone_type="anchor",
            function="Grips left hip/pelvis to stabilize brace and prevent riding up",
            direction="grip",
            depth_mm=4.0,
            radius_mm=(40.0, 60.0, 40.0)
        ),
        PressureZone(
            name="Right Hip Anchor",
            marker_name="LM_ANCHOR_HIP_R",
            position=tuple(marker_pos.get("LM_ANCHOR_HIP_R", [0, 0, 0])),
            zone_type="anchor",
            function="Grips right hip/pelvis to stabilize brace and prevent riding up",
            direction="grip",
            depth_mm=4.0,
            radius_mm=(40.0, 60.0, 40.0)
        ),
    ]
    
    return zones


def transform_markers(
    markers: Dict[str, list],
    transform_matrix: np.ndarray
) -> Dict[str, Tuple[float, float, float]]:
    """Apply transformation matrix to all marker positions."""
    transformed = {}
    
    for name, pos in markers.items():
        if isinstance(pos, (list, tuple)) and len(pos) == 3:
            # Convert to homogeneous coordinates
            pos_h = np.array([pos[0], pos[1], pos[2], 1.0])
            # Apply transform
            new_pos = transform_matrix @ pos_h
            transformed[name] = (float(new_pos[0]), float(new_pos[1]), float(new_pos[2]))
    
    return transformed


def generate_glb_brace(
    rigo_type: str,
    template_type: TemplateType,
    output_dir: Path,
    case_id: str,
    cobb_angles: Dict[str, float],
    body_scan_path: Optional[str] = None,
    clearance_mm: float = 8.0
) -> BraceGenerationResult:
    """
    Generate a GLB brace with markers.
    
    Args:
        rigo_type: Rigo classification (A1, A2, etc.)
        template_type: "regular" or "vase"
        output_dir: Directory for output files
        case_id: Case identifier
        cobb_angles: Dict with PT, MT, TL angles
        body_scan_path: Optional path to body scan STL for fitting
        clearance_mm: Clearance between body and brace
        
    Returns:
        BraceGenerationResult with paths and marker info
    """
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Load template
    mesh = load_glb_template(rigo_type, template_type)
    marker_data = load_template_markers(rigo_type, template_type)
    
    markers = marker_data.get("markers", {})
    basis = marker_data.get("basis", {})
    
    transform_matrix = np.eye(4)
    transform_info = None
    
    # If body scan provided, fit to body
    if body_scan_path and Path(body_scan_path).exists():
        mesh, transform_matrix, transform_info = fit_brace_to_body(
            mesh, body_scan_path, clearance_mm, brace_basis=basis, template_type=template_type,
            markers=markers  # Pass markers for zone-aware ironing
        )
        # Transform markers to match the body-fitted mesh
        markers = transform_markers(markers, transform_matrix)
    
    # Calculate pressure zones
    pressure_zones = calculate_pressure_zones(markers, rigo_type, cobb_angles)
    
    # Output file names
    type_suffix = "vase" if template_type == "vase" else "regular"
    glb_filename = f"{case_id}_{rigo_type}_{type_suffix}.glb"
    stl_filename = f"{case_id}_{rigo_type}_{type_suffix}.stl"
    json_filename = f"{case_id}_{rigo_type}_{type_suffix}_markers.json"
    
    glb_path = output_dir / glb_filename
    stl_path = output_dir / stl_filename
    json_path = output_dir / json_filename
    
    # Export GLB
    mesh.export(str(glb_path))
    
    # Export STL
    mesh.export(str(stl_path))
    
    # Build output JSON with markers and zones
    output_data = {
        "case_id": case_id,
        "rigo_type": rigo_type,
        "template_type": template_type,
        "cobb_angles": cobb_angles,
        "markers": {k: list(v) if isinstance(v, tuple) else v for k, v in markers.items()},
        "basis": basis,
        "pressure_zones": [
            {
                "name": z.name,
                "marker_name": z.marker_name,
                "position": list(z.position),
                "zone_type": z.zone_type,
                "function": z.function,
                "direction": z.direction,
                "depth_mm": z.depth_mm,
                "radius_mm": list(z.radius_mm)
            }
            for z in pressure_zones
        ],
        "mesh_stats": {
            "vertices": len(mesh.vertices),
            "faces": len(mesh.faces)
        },
        "outputs": {
            "glb": str(glb_path),
            "stl": str(stl_path),
            "json": str(json_path)
        }
    }
    
    if transform_info:
        output_data["body_fitting"] = transform_info
    
    # Save JSON
    with open(json_path, "w") as f:
        json.dump(output_data, f, indent=2)
    
    return BraceGenerationResult(
        glb_path=str(glb_path),
        stl_path=str(stl_path),
        json_path=str(json_path),
        template_type=template_type,
        rigo_type=rigo_type,
        markers=markers,
        basis=basis,
        pressure_zones=[asdict(z) for z in pressure_zones],
        mesh_stats=output_data["mesh_stats"],
        transform_applied=transform_info
    )


def iron_brace_to_body(
    brace_mesh: trimesh.Trimesh,
    body_mesh: trimesh.Trimesh,
    min_clearance_mm: float = 3.0,
    max_clearance_mm: float = 15.0,
    smoothing_iterations: int = 2,
    up_axis: int = 2,
    markers: Optional[Dict[str, Any]] = None
) -> trimesh.Trimesh:
    """
    Iron the brace surface to conform to the body scan surface.
    
    This ensures the brace follows the body contour without excessive gaps.
    Uses zone-aware ironing:
    - FRONT (belly) and BACK: Aggressive ironing for tight fit
    - SIDES (where pads/bays are): Preserve correction zones, moderate ironing
    
    Args:
        brace_mesh: The brace mesh to iron
        body_mesh: The body scan mesh to conform to
        min_clearance_mm: Minimum distance from body surface
        max_clearance_mm: Maximum distance from body surface (trigger ironing)
        smoothing_iterations: Number of Laplacian smoothing passes after ironing
        up_axis: Which axis is "up" (0=X, 1=Y, 2=Z)
        markers: Optional dict of marker positions to preserve pressure zones
    
    Returns:
        Ironed brace mesh
    """
    from scipy.spatial import cKDTree
    import math
    
    print(f"Ironing brace to body surface (clearance: {min_clearance_mm}-{max_clearance_mm}mm)")
    
    # Create a copy to modify
    ironed_mesh = brace_mesh.copy()
    vertices = ironed_mesh.vertices.copy()
    
    # Get body center and bounds
    body_center = body_mesh.centroid
    body_bounds = body_mesh.bounds
    
    # Determine the torso region (process middle 80% of body height)
    body_height = body_bounds[1, up_axis] - body_bounds[0, up_axis]
    torso_bottom = body_bounds[0, up_axis] + body_height * 0.10
    torso_top = body_bounds[0, up_axis] + body_height * 0.90
    
    # Build KD-tree from body mesh vertices for fast nearest neighbor queries
    body_tree = cKDTree(body_mesh.vertices)
    
    # Find closest points on body for ALL brace vertices at once
    distances, closest_indices = body_tree.query(vertices, k=1)
    closest_points = body_mesh.vertices[closest_indices]
    
    # Determine horizontal axes (perpendicular to up axis)
    horiz_axes = [i for i in range(3) if i != up_axis]
    
    # Calculate brace center for angle computation
    brace_center = np.mean(vertices, axis=0)
    
    # Identify marker exclusion zones (preserve correction areas)
    exclusion_zones = []
    if markers:
        # Pad and bay markers need preservation
        for marker_name in ['LM_PAD_TH', 'LM_PAD_LUM', 'LM_BAY_TH', 'LM_BAY_LUM']:
            if marker_name in markers:
                pos = markers[marker_name]
                if isinstance(pos, (list, tuple)) and len(pos) >= 3:
                    exclusion_zones.append({
                        'center': np.array(pos),
                        'radius': 60.0,  # 60mm exclusion radius around markers
                        'name': marker_name
                    })
    
    # Process each brace vertex
    adjusted_count = 0
    pulled_in_count = 0
    pushed_out_count = 0
    skipped_zone_count = 0
    
    # Height normalization
    brace_min_z = vertices[:, up_axis].min()
    brace_max_z = vertices[:, up_axis].max()
    brace_height_range = max(brace_max_z - brace_min_z, 1.0)
    
    for i in range(len(vertices)):
        vertex = vertices[i]
        closest_pt = closest_points[i]
        dist = distances[i]
        
        # Only process vertices in the torso region
        if vertex[up_axis] < torso_bottom or vertex[up_axis] > torso_top:
            continue
        
        # Check if vertex is in an exclusion zone (near pad/bay markers)
        in_exclusion = False
        for zone in exclusion_zones:
            zone_dist = np.linalg.norm(vertex - zone['center'])
            if zone_dist < zone['radius']:
                in_exclusion = True
                skipped_zone_count += 1
                break
        
        if in_exclusion:
            continue
        
        # Calculate angular position around body center (horizontal plane)
        # 0° = front (belly), 90° = right side, 180° = back, 270° = left side
        rel_pos = vertex - body_center
        angle = math.atan2(rel_pos[horiz_axes[1]], rel_pos[horiz_axes[0]])
        angle_deg = math.degrees(angle) % 360
        
        # Determine zone based on angle:
        # FRONT (belly): 315-45° - aggressive ironing
        # BACK: 135-225° - aggressive ironing  
        # SIDES: 45-135° and 225-315° - moderate ironing (correction zones)
        is_front_back = (angle_deg < 45 or angle_deg > 315) or (135 < angle_deg < 225)
        
        # Height-based clearance adjustment
        height_norm = (vertex[up_axis] - brace_min_z) / brace_height_range
        
        # Set clearances based on zone
        if is_front_back:
            # FRONT/BACK: Aggressive ironing - very tight fit
            local_min = min_clearance_mm * 0.5  # Allow closer to body
            local_max = max_clearance_mm * 0.6  # Trigger ironing earlier
            local_target = min_clearance_mm + 2.0  # Target just above minimum
        else:
            # SIDES: More conservative - preserve room for correction
            local_min = min_clearance_mm
            local_max = max_clearance_mm * 1.2  # Allow slightly more gap
            local_target = (min_clearance_mm + max_clearance_mm) / 2
        
        # Height adjustments (tighter at hips and chest)
        if height_norm < 0.25 or height_norm > 0.75:
            local_max *= 0.8  # Tighter at extremes
            local_target *= 0.85
        
        # Direction from body surface to brace vertex
        direction = vertex - closest_pt
        dir_length = np.linalg.norm(direction)
        
        if dir_length < 1e-6:
            direction = vertex - body_center
            direction[up_axis] = 0
            dir_length = np.linalg.norm(direction)
            if dir_length < 1e-6:
                continue
        
        direction = direction / dir_length
        
        # Determine signed distance
        vertex_dist_to_center = np.linalg.norm(vertex[:2] - body_center[:2])
        closest_dist_to_center = np.linalg.norm(closest_pt[:2] - body_center[:2])
        
        if vertex_dist_to_center >= closest_dist_to_center:
            signed_distance = dist
        else:
            signed_distance = -dist
        
        # Determine if adjustment is needed
        needs_adjustment = False
        new_position = vertex.copy()
        
        if signed_distance > local_max:
            # Gap too large - pull vertex closer to body
            new_position = closest_pt + direction * local_target
            new_position[up_axis] = vertex[up_axis]  # Preserve height
            needs_adjustment = True
            pulled_in_count += 1
            
        elif signed_distance < local_min:
            # Too close or inside body - push outward
            offset = local_min + 1.0
            outward_dir = closest_pt - body_center
            outward_dir[up_axis] = 0
            outward_length = np.linalg.norm(outward_dir)
            if outward_length > 1e-6:
                outward_dir = outward_dir / outward_length
                new_position = closest_pt + outward_dir * offset
                new_position[up_axis] = vertex[up_axis]
                needs_adjustment = True
                pushed_out_count += 1
        
        if needs_adjustment:
            vertices[i] = new_position
            adjusted_count += 1
    
    print(f"Ironing adjusted {adjusted_count} vertices (pulled in: {pulled_in_count}, pushed out: {pushed_out_count}, skipped zones: {skipped_zone_count})")
    
    # Apply modified vertices
    ironed_mesh.vertices = vertices
    
    # Apply Laplacian smoothing to blend changes and remove artifacts
    if smoothing_iterations > 0 and adjusted_count > 0:
        print(f"Applying {smoothing_iterations} smoothing iterations")
        try:
            ironed_mesh = trimesh.smoothing.filter_laplacian(
                ironed_mesh, 
                lamb=0.3,  # Gentler smoothing to preserve shape
                iterations=smoothing_iterations,
                implicit_time_integration=False
            )
        except Exception as e:
            print(f"Smoothing failed (non-critical): {e}")
    
    # Ensure mesh is valid
    ironed_mesh.fix_normals()
    
    return ironed_mesh


def fit_brace_to_body(
    brace_mesh: trimesh.Trimesh,
    body_scan_path: str,
    clearance_mm: float = 8.0,
    brace_basis: Optional[Dict[str, Any]] = None,
    template_type: str = "regular",
    enable_ironing: bool = True,
    markers: Optional[Dict[str, Any]] = None
) -> Tuple[trimesh.Trimesh, np.ndarray, Dict[str, Any]]:
    """
    Fit brace to body scan using basis alignment.
    
    The brace needs to be:
    1. Rotated so its UP axis aligns with body's UP axis (typically Z for body scans)
    2. Scaled to fit around the body with proper clearance
    3. Positioned at the torso level
    4. Ironed to conform to body surface (respecting correction zones)
    
    Returns:
        Tuple of (transformed_mesh, transform_matrix, fitting_info)
    """
    # Load body scan
    body_mesh = trimesh.load(body_scan_path, force='mesh')
    
    # Get body dimensions
    body_bounds = body_mesh.bounds
    body_extents = body_mesh.extents
    body_center = body_mesh.centroid
    
    # Determine body up axis (typically the longest dimension = height)
    # For human body scans, this is usually Z (from 3D scanners) or Y
    body_up_axis_idx = np.argmax(body_extents)
    print(f"Body up axis: {['X', 'Y', 'Z'][body_up_axis_idx]}, extents: {body_extents}")
    
    # Get brace dimensions
    brace_bounds = brace_mesh.bounds
    brace_extents = brace_mesh.extents
    brace_center = brace_mesh.centroid
    
    print(f"Brace original extents: {brace_extents}, template_type: {template_type}")
    
    # Start building transformation
    transformed_mesh = brace_mesh.copy()
    transform = np.eye(4)
    
    # Step 1: Center brace at origin
    T_center = np.eye(4)
    T_center[:3, 3] = -brace_center
    transformed_mesh.apply_transform(T_center)
    transform = T_center @ transform
    
    # Step 2: Apply rotations based on template type and body orientation
    # Regular templates have: negative Y is up (inverted), need to flip
    # Vase templates have: positive Y is up
    # Body scan is Z-up
    
    if body_up_axis_idx == 2:  # Body is Z-up (standard for 3D scanners)
        if template_type == "regular":
            # Regular brace: -Y is up (inverted)
            # 1. Rotate -90° around X to bring Y-up to Z-up
            R1 = trimesh.transformations.rotation_matrix(-np.pi/2, [1, 0, 0])
            transformed_mesh.apply_transform(R1)
            transform = R1 @ transform
            
            # 2. The brace is now Z-up but inverted (pelvis at top, shoulders at bottom)
            #    Flip 180° around X to correct (this keeps Z as up axis)
            R2 = trimesh.transformations.rotation_matrix(np.pi, [1, 0, 0])
            transformed_mesh.apply_transform(R2)
            transform = R2 @ transform
            
            # 3. Rotate around Z to face forward correctly
            R3 = trimesh.transformations.rotation_matrix(-np.pi/2, [0, 0, 1])
            transformed_mesh.apply_transform(R3)
            transform = R3 @ transform
            
            print(f"Applied regular brace rotations: X-90°, X+180° (flip), Z-90°")
            
        else:  # vase
            # Vase brace: positive Y is up
            # 1. Rotate -90° around X to bring Y-up to Z-up
            R1 = trimesh.transformations.rotation_matrix(-np.pi/2, [1, 0, 0])
            transformed_mesh.apply_transform(R1)
            transform = R1 @ transform
            
            # 2. Flip 180° around Y to correct orientation (right-side up)
            R2 = trimesh.transformations.rotation_matrix(np.pi, [0, 1, 0])
            transformed_mesh.apply_transform(R2)
            transform = R2 @ transform
            
            print(f"Applied vase brace rotations: X-90°, Y+180° (flip)")
    
    # Step 3: Get new brace dimensions after rotation
    new_brace_extents = transformed_mesh.extents
    new_brace_center = transformed_mesh.centroid
    print(f"Brace extents after rotation: {new_brace_extents}")
    
    # Step 4: Calculate NON-UNIFORM scaling based on body dimensions
    # The brace should cover the TORSO region (~50% of body height)
    # AND wrap around the body with proper girth
    
    body_height = body_extents[body_up_axis_idx]
    brace_height = new_brace_extents[body_up_axis_idx]  # After rotation, this is the height
    
    # Body horizontal dimensions (girth at torso level)
    horizontal_axes = [i for i in range(3) if i != body_up_axis_idx]
    body_width = body_extents[horizontal_axes[0]]  # X width
    body_depth = body_extents[horizontal_axes[1]]  # Y depth
    
    # Brace horizontal dimensions
    brace_width = new_brace_extents[horizontal_axes[0]]
    brace_depth = new_brace_extents[horizontal_axes[1]]
    
    # Target: brace height should cover ~65% of body height (full torso coverage)
    target_height = body_height * 0.65
    height_scale = target_height / brace_height if brace_height > 0 else 1.0
    
    # Target: brace width/depth should be LARGER than body to wrap AROUND it
    # The brace sits OUTSIDE the body, only pressure points push inward
    # Add ~25% extra + clearance so brace externals are visible outside body
    target_width = body_width * 1.25 + clearance_mm * 2
    target_depth = body_depth * 1.25 + clearance_mm * 2
    
    width_scale = target_width / brace_width if brace_width > 0 else 1.0
    depth_scale = target_depth / brace_depth if brace_depth > 0 else 1.0
    
    # Apply non-uniform scaling
    # Determine which axis is which after rotation
    S = np.eye(4)
    if body_up_axis_idx == 2:  # Z is up
        S[0, 0] = width_scale   # X scale
        S[1, 1] = depth_scale   # Y scale  
        S[2, 2] = height_scale  # Z scale (height)
    elif body_up_axis_idx == 1:  # Y is up
        S[0, 0] = width_scale   # X scale
        S[1, 1] = height_scale  # Y scale (height)
        S[2, 2] = depth_scale   # Z scale
    else:  # X is up (unusual)
        S[0, 0] = height_scale  # X scale (height)
        S[1, 1] = width_scale   # Y scale
        S[2, 2] = depth_scale   # Z scale
    
    # Limit scales to reasonable range
    S[0, 0] = max(0.5, min(S[0, 0], 50.0))
    S[1, 1] = max(0.5, min(S[1, 1], 50.0))
    S[2, 2] = max(0.5, min(S[2, 2], 50.0))
    
    transformed_mesh.apply_transform(S)
    transform = S @ transform
    
    print(f"Applied non-uniform scale: width={S[0,0]:.2f}, depth={S[1,1]:.2f}, height={S[2,2]:.2f}")
    print(f"Target dimensions: width={target_width:.1f}, depth={target_depth:.1f}, height={target_height:.1f}")
    
    # For fitting_info, use average scale
    scale = (S[0, 0] + S[1, 1] + S[2, 2]) / 3
    
    # Step 6: Position brace at torso level
    # Calculate where the torso is (middle portion of body height)
    body_height = body_extents[body_up_axis_idx]
    body_bottom = body_bounds[0, body_up_axis_idx]
    body_top = body_bounds[1, body_up_axis_idx]
    
    # Torso is roughly the middle 40% of body height (from ~30% to ~70%)
    torso_center_ratio = 0.5  # Middle of body
    torso_center_height = body_bottom + body_height * torso_center_ratio
    
    # Target position: center horizontally on body, at torso height vertically
    target_center = body_center.copy()
    target_center[body_up_axis_idx] = torso_center_height
    
    # Current brace center after transformations
    current_center = transformed_mesh.centroid
    
    T_position = np.eye(4)
    T_position[:3, 3] = target_center - current_center
    transformed_mesh.apply_transform(T_position)
    transform = T_position @ transform
    
    # Step 7: Iron brace to conform to body surface (eliminate gaps and humps)
    # Transform markers so we can exclude correction zones from ironing
    transformed_markers = None
    if markers:
        transformed_markers = transform_markers(markers, transform)
    
    ironing_info = {}
    if enable_ironing:
        try:
            print(f"Starting brace ironing to body surface...")
            pre_iron_extents = transformed_mesh.extents.copy()
            
            transformed_mesh = iron_brace_to_body(
                brace_mesh=transformed_mesh,
                body_mesh=body_mesh,
                min_clearance_mm=clearance_mm * 0.4,  # Allow closer for tight fit
                max_clearance_mm=clearance_mm * 1.5,  # Iron areas with gaps > 1.5x clearance
                smoothing_iterations=3,
                up_axis=body_up_axis_idx,
                markers=transformed_markers
            )
            
            post_iron_extents = transformed_mesh.extents
            ironing_info = {
                "enabled": True,
                "pre_iron_extents": pre_iron_extents.tolist(),
                "post_iron_extents": post_iron_extents.tolist(),
                "min_clearance_mm": clearance_mm * 0.5,
                "max_clearance_mm": clearance_mm * 2.0,
            }
            print(f"Ironing complete. Extents changed from {pre_iron_extents} to {post_iron_extents}")
            
        except Exception as e:
            print(f"Ironing failed (non-critical): {e}")
            ironing_info = {"enabled": False, "error": str(e)}
    else:
        ironing_info = {"enabled": False}
    
    fitting_info = {
        "scale_avg": float(scale),
        "scale_x": float(S[0, 0]),
        "scale_y": float(S[1, 1]),
        "scale_z": float(S[2, 2]),
        "template_type": template_type,
        "body_extents": body_extents.tolist(),
        "brace_extents_original": brace_extents.tolist(),
        "brace_extents_final": transformed_mesh.extents.tolist(),
        "clearance_mm": clearance_mm,
        "body_center": body_center.tolist(),
        "final_center": transformed_mesh.centroid.tolist(),
        "body_up_axis": int(body_up_axis_idx),
        "ironing": ironing_info,
    }
    
    return transformed_mesh, transform, fitting_info


def generate_both_brace_types(
    rigo_type: str,
    output_dir: Path,
    case_id: str,
    cobb_angles: Dict[str, float],
    body_scan_path: Optional[str] = None,
    clearance_mm: float = 8.0
) -> Dict[str, BraceGenerationResult]:
    """
    Generate both regular and vase brace types for comparison.
    
    Returns:
        Dict with "regular" and "vase" results
    """
    results = {}
    
    # Generate regular brace
    try:
        results["regular"] = generate_glb_brace(
            rigo_type=rigo_type,
            template_type="regular",
            output_dir=output_dir,
            case_id=case_id,
            cobb_angles=cobb_angles,
            body_scan_path=body_scan_path,
            clearance_mm=clearance_mm
        )
    except FileNotFoundError as e:
        results["regular"] = {"error": str(e)}
    
    # Generate vase brace
    try:
        results["vase"] = generate_glb_brace(
            rigo_type=rigo_type,
            template_type="vase",
            output_dir=output_dir,
            case_id=case_id,
            cobb_angles=cobb_angles,
            body_scan_path=body_scan_path,
            clearance_mm=clearance_mm
        )
    except FileNotFoundError as e:
        results["vase"] = {"error": str(e)}
    
    return results


# Available templates
AVAILABLE_RIGO_TYPES = ["A1", "A2", "A3", "B1", "B2", "C1", "C2", "E1", "E2"]

def list_available_templates() -> Dict[str, list]:
    """List all available template files."""
    regular = []
    vase = []
    
    for rigo_type in AVAILABLE_RIGO_TYPES:
        glb_path, _ = get_template_paths(rigo_type, "regular")
        if glb_path.exists():
            regular.append(rigo_type)
        
        glb_path, _ = get_template_paths(rigo_type, "vase")
        if glb_path.exists():
            vase.append(rigo_type)
    
    return {
        "regular": regular,
        "vase": vase
    }
