models/rigging_processor.py

import numpy as np
import trimesh
from typing import Union, Dict, List, Tuple, Optional
import tempfile
from pathlib import Path

class UniRigProcessor:
    """Automatic rigging for 3D models using simplified UniRig approach"""
    
    def __init__(self, device: str = "cuda"):
        self.device = device
        self.model = None
        
        # Rigging parameters
        self.bone_detection_threshold = 0.1
        self.max_bones = 20
        self.min_bones = 5
        
        # Animation presets for monsters
        self.animation_presets = {
            'idle': self._create_idle_animation,
            'walk': self._create_walk_animation,
            'attack': self._create_attack_animation,
            'happy': self._create_happy_animation
        }
    
    def load_model(self):
        """Load rigging model (placeholder for actual implementation)"""
        # In production, this would load the actual UniRig model
        # For now, we'll use procedural rigging
        self.model = "procedural"
    
    def rig_mesh(self, 
                mesh: Union[str, trimesh.Trimesh],
                mesh_type: str = "monster") -> Dict[str, any]:
        """Add rigging to a 3D mesh"""
        
        try:
            # Load mesh if path provided
            if isinstance(mesh, str):
                mesh = trimesh.load(mesh)
            
            # Ensure model is loaded
            if self.model is None:
                self.load_model()
            
            # Analyze mesh structure
            mesh_analysis = self._analyze_mesh(mesh)
            
            # Generate skeleton
            skeleton = self._generate_skeleton(mesh, mesh_analysis)
            
            # Compute bone weights
            weights = self._compute_bone_weights(mesh, skeleton)
            
            # Create rigged model
            rigged_model = {
                'mesh': mesh,
                'skeleton': skeleton,
                'weights': weights,
                'animations': self._create_default_animations(skeleton),
                'metadata': {
                    'mesh_type': mesh_type,
                    'bone_count': len(skeleton['bones']),
                    'vertex_count': len(mesh.vertices)
                }
            }
            
            # Save rigged model
            output_path = self._save_rigged_model(rigged_model)
            
            return output_path
            
        except Exception as e:
            print(f"Rigging error: {e}")
            # Return original mesh if rigging fails
            return self._save_mesh_without_rigging(mesh)
    
    def _analyze_mesh(self, mesh: trimesh.Trimesh) -> Dict[str, any]:
        """Analyze mesh structure for rigging"""
        
        # Get mesh bounds and center
        bounds = mesh.bounds
        center = mesh.centroid
        
        # Analyze mesh topology
        analysis = {
            'bounds': bounds,
            'center': center,
            'height': bounds[1][2] - bounds[0][2],
            'width': bounds[1][0] - bounds[0][0],
            'depth': bounds[1][1] - bounds[0][1],
            'is_symmetric': self._check_symmetry(mesh),
            'detected_limbs': self._detect_limbs(mesh),
            'mesh_type': self._classify_mesh_type(mesh)
        }
        
        return analysis
    
    def _check_symmetry(self, mesh: trimesh.Trimesh) -> bool:
        """Check if mesh is roughly symmetric"""
        # Simple check: compare left and right halves
        vertices = mesh.vertices
        center_x = mesh.centroid[0]
        
        left_verts = vertices[vertices[:, 0] < center_x]
        right_verts = vertices[vertices[:, 0] > center_x]
        
        # Check if similar number of vertices on each side
        ratio = len(left_verts) / (len(right_verts) + 1)
        return 0.8 < ratio < 1.2
    
    def _detect_limbs(self, mesh: trimesh.Trimesh) -> List[Dict]:
        """Detect potential limbs in the mesh"""
        # Simplified limb detection using vertex clustering
        from sklearn.cluster import DBSCAN
        
        limbs = []
        
        try:
            # Cluster vertices to find distinct parts
            clustering = DBSCAN(eps=0.1, min_samples=10).fit(mesh.vertices)
            
            # Analyze each cluster
            for label in set(clustering.labels_):
                if label == -1:  # Noise
                    continue
                
                cluster_verts = mesh.vertices[clustering.labels_ == label]
                
                # Check if cluster could be a limb
                cluster_bounds = np.array([cluster_verts.min(axis=0), cluster_verts.max(axis=0)])
                dimensions = cluster_bounds[1] - cluster_bounds[0]
                
                # Limbs are typically elongated
                if max(dimensions) / (min(dimensions) + 0.001) > 2:
                    limbs.append({
                        'center': cluster_verts.mean(axis=0),
                        'direction': dimensions,
                        'size': len(cluster_verts)
                    })
        except:
            # Fallback if clustering fails
            pass
        
        return limbs
    
    def _classify_mesh_type(self, mesh: trimesh.Trimesh) -> str:
        """Classify the type of creature mesh"""
        analysis = {
            'height': mesh.bounds[1][2] - mesh.bounds[0][2],
            'width': mesh.bounds[1][0] - mesh.bounds[0][0],
            'depth': mesh.bounds[1][1] - mesh.bounds[0][1]
        }
        
        # Simple classification based on proportions
        aspect_ratio = analysis['height'] / max(analysis['width'], analysis['depth'])
        
        if aspect_ratio > 1.5:
            return 'bipedal'  # Tall creatures
        elif aspect_ratio < 0.7:
            return 'quadruped'  # Wide creatures
        else:
            return 'hybrid'  # Mixed proportions
    
    def _generate_skeleton(self, mesh: trimesh.Trimesh, analysis: Dict) -> Dict:
        """Generate skeleton for the mesh"""
        
        skeleton = {
            'bones': [],
            'hierarchy': {},
            'bind_poses': []
        }
        
        # Create root bone at center
        root_pos = analysis['center']
        root_bone = {
            'id': 0,
            'name': 'root',
            'position': root_pos,
            'parent': -1,
            'children': []
        }
        skeleton['bones'].append(root_bone)
        
        # Generate bones based on mesh type
        mesh_type = analysis['mesh_type']
        
        if mesh_type == 'bipedal':
            skeleton = self._generate_bipedal_skeleton(mesh, skeleton, analysis)
        elif mesh_type == 'quadruped':
            skeleton = self._generate_quadruped_skeleton(mesh, skeleton, analysis)
        else:
            skeleton = self._generate_hybrid_skeleton(mesh, skeleton, analysis)
        
        # Build hierarchy
        for bone in skeleton['bones']:
            if bone['parent'] >= 0:
                skeleton['bones'][bone['parent']]['children'].append(bone['id'])
        
        return skeleton
    
    def _generate_bipedal_skeleton(self, mesh: trimesh.Trimesh, skeleton: Dict, analysis: Dict) -> Dict:
        """Generate skeleton for bipedal creature"""
        
        bounds = analysis['bounds']
        center = analysis['center']
        height = analysis['height']
        
        # Spine bones
        spine_positions = [
            center + [0, 0, -height * 0.4],  # Hips
            center + [0, 0, 0],               # Chest
            center + [0, 0, height * 0.3]     # Head
        ]
        
        parent_id = 0
        for i, pos in enumerate(spine_positions):
            bone = {
                'id': len(skeleton['bones']),
                'name': ['hips', 'chest', 'head'][i],
                'position': pos,
                'parent': parent_id,
                'children': []
            }
            skeleton['bones'].append(bone)
            parent_id = bone['id']
        
        # Add limbs
        chest_id = skeleton['bones'][2]['id']  # Chest bone
        hips_id = skeleton['bones'][1]['id']   # Hips bone
        
        # Arms
        arm_offset = analysis['width'] * 0.4
        for side, offset in [('left', -arm_offset), ('right', arm_offset)]:
            shoulder_pos = skeleton['bones'][chest_id]['position'] + [offset, 0, 0]
            elbow_pos = shoulder_pos + [offset * 0.5, 0, -height * 0.2]
            
            # Shoulder
            shoulder = {
                'id': len(skeleton['bones']),
                'name': f'{side}_shoulder',
                'position': shoulder_pos,
                'parent': chest_id,
                'children': []
            }
            skeleton['bones'].append(shoulder)
            
            # Elbow/Hand
            hand = {
                'id': len(skeleton['bones']),
                'name': f'{side}_hand',
                'position': elbow_pos,
                'parent': shoulder['id'],
                'children': []
            }
            skeleton['bones'].append(hand)
        
        # Legs
        for side, offset in [('left', -arm_offset * 0.5), ('right', arm_offset * 0.5)]:
            hip_pos = skeleton['bones'][hips_id]['position'] + [offset, 0, 0]
            foot_pos = hip_pos + [0, 0, -height * 0.4]
            
            # Leg
            leg = {
                'id': len(skeleton['bones']),
                'name': f'{side}_leg',
                'position': hip_pos,
                'parent': hips_id,
                'children': []
            }
            skeleton['bones'].append(leg)
            
            # Foot
            foot = {
                'id': len(skeleton['bones']),
                'name': f'{side}_foot',
                'position': foot_pos,
                'parent': leg['id'],
                'children': []
            }
            skeleton['bones'].append(foot)
        
        return skeleton
    
    def _generate_quadruped_skeleton(self, mesh: trimesh.Trimesh, skeleton: Dict, analysis: Dict) -> Dict:
        """Generate skeleton for quadruped creature"""
        
        # Similar to bipedal but with 4 legs and horizontal spine
        center = analysis['center']
        width = analysis['width']
        depth = analysis['depth']
        
        # Spine (horizontal)
        spine_positions = [
            center + [-width * 0.3, 0, 0],  # Tail
            center,                          # Body
            center + [width * 0.3, 0, 0]    # Head
        ]
        
        parent_id = 0
        for i, pos in enumerate(spine_positions):
            bone = {
                'id': len(skeleton['bones']),
                'name': ['tail', 'body', 'head'][i],
                'position': pos,
                'parent': parent_id,
                'children': []
            }
            skeleton['bones'].append(bone)
            parent_id = bone['id'] if i < 2 else skeleton['bones'][1]['id']
        
        # Add 4 legs
        body_id = skeleton['bones'][1]['id']
        
        for front_back, x_offset in [('front', width * 0.2), ('back', -width * 0.2)]:
            for side, z_offset in [('left', -depth * 0.3), ('right', depth * 0.3)]:
                leg_pos = skeleton['bones'][body_id]['position'] + [x_offset, -analysis['height'] * 0.3, z_offset]
                
                leg = {
                    'id': len(skeleton['bones']),
                    'name': f'{front_back}_{side}_leg',
                    'position': leg_pos,
                    'parent': body_id,
                    'children': []
                }
                skeleton['bones'].append(leg)
        
        return skeleton
    
    def _generate_hybrid_skeleton(self, mesh: trimesh.Trimesh, skeleton: Dict, analysis: Dict) -> Dict:
        """Generate skeleton for hybrid creature"""
        # Mix of bipedal and quadruped features
        # For simplicity, use bipedal as base
        return self._generate_bipedal_skeleton(mesh, skeleton, analysis)
    
    def _compute_bone_weights(self, mesh: trimesh.Trimesh, skeleton: Dict) -> np.ndarray:
        """Compute bone weights for vertices"""
        
        num_vertices = len(mesh.vertices)
        num_bones = len(skeleton['bones'])
        
        # Initialize weights matrix
        weights = np.zeros((num_vertices, num_bones))
        
        # For each vertex, compute influence from each bone
        for v_idx, vertex in enumerate(mesh.vertices):
            total_weight = 0
            
            for b_idx, bone in enumerate(skeleton['bones']):
                # Distance-based weight
                distance = np.linalg.norm(vertex - bone['position'])
                
                # Inverse distance weight with falloff
                weight = 1.0 / (distance + 0.1)
                weights[v_idx, b_idx] = weight
                total_weight += weight
            
            # Normalize weights
            if total_weight > 0:
                weights[v_idx] /= total_weight
            
            # Keep only top 4 influences per vertex (standard for game engines)
            top_4 = np.argsort(weights[v_idx])[-4:]
            mask = np.zeros(num_bones, dtype=bool)
            mask[top_4] = True
            weights[v_idx, ~mask] = 0
            
            # Re-normalize
            if weights[v_idx].sum() > 0:
                weights[v_idx] /= weights[v_idx].sum()
        
        return weights
    
    def _create_default_animations(self, skeleton: Dict) -> Dict[str, List]:
        """Create default animations for the skeleton"""
        
        animations = {}
        
        # Create basic animation sets
        for anim_name, anim_func in self.animation_presets.items():
            animations[anim_name] = anim_func(skeleton)
        
        return animations
    
    def _create_idle_animation(self, skeleton: Dict) -> List[Dict]:
        """Create idle animation keyframes"""
        
        keyframes = []
        
        # Simple breathing/bobbing motion
        for t in np.linspace(0, 2 * np.pi, 30):
            frame = {
                'time': t / (2 * np.pi),
                'bones': {}
            }
            
            # Subtle movement for each bone
            for bone in skeleton['bones']:
                if 'chest' in bone['name'] or 'body' in bone['name']:
                    # Breathing motion
                    offset = np.sin(t) * 0.02
                    frame['bones'][bone['id']] = {
                        'position': bone['position'] + [0, offset, 0],
                        'rotation': [0, 0, 0, 1]  # Quaternion
                    }
                else:
                    # No movement
                    frame['bones'][bone['id']] = {
                        'position': bone['position'],
                        'rotation': [0, 0, 0, 1]
                    }
            
            keyframes.append(frame)
        
        return keyframes
    
    def _create_walk_animation(self, skeleton: Dict) -> List[Dict]:
        """Create walk animation keyframes"""
        # Simplified walk cycle
        keyframes = []
        
        for t in np.linspace(0, 2 * np.pi, 60):
            frame = {
                'time': t / (2 * np.pi),
                'bones': {}
            }
            
            # Animate legs with sine waves
            for bone in skeleton['bones']:
                if 'leg' in bone['name'] or 'foot' in bone['name']:
                    # Alternating leg movement
                    phase = 0 if 'left' in bone['name'] else np.pi
                    offset = np.sin(t + phase) * 0.1
                    
                    frame['bones'][bone['id']] = {
                        'position': bone['position'] + [offset, 0, 0],
                        'rotation': [0, 0, 0, 1]
                    }
                else:
                    frame['bones'][bone['id']] = {
                        'position': bone['position'],
                        'rotation': [0, 0, 0, 1]
                    }
            
            keyframes.append(frame)
        
        return keyframes
    
    def _create_attack_animation(self, skeleton: Dict) -> List[Dict]:
        """Create attack animation keyframes"""
        # Quick strike motion
        keyframes = []
        
        # Wind up
        for t in np.linspace(0, 0.3, 10):
            frame = {'time': t, 'bones': {}}
            for bone in skeleton['bones']:
                frame['bones'][bone['id']] = {
                    'position': bone['position'],
                    'rotation': [0, 0, 0, 1]
                }
            keyframes.append(frame)
        
        # Strike
        for t in np.linspace(0.3, 0.5, 5):
            frame = {'time': t, 'bones': {}}
            for bone in skeleton['bones']:
                if 'hand' in bone['name'] or 'head' in bone['name']:
                    # Forward motion
                    offset = (t - 0.3) * 0.5
                    frame['bones'][bone['id']] = {
                        'position': bone['position'] + [offset, 0, 0],
                        'rotation': [0, 0, 0, 1]
                    }
                else:
                    frame['bones'][bone['id']] = {
                        'position': bone['position'],
                        'rotation': [0, 0, 0, 1]
                    }
            keyframes.append(frame)
        
        # Return
        for t in np.linspace(0.5, 1.0, 10):
            frame = {'time': t, 'bones': {}}
            for bone in skeleton['bones']:
                frame['bones'][bone['id']] = {
                    'position': bone['position'],
                    'rotation': [0, 0, 0, 1]
                }
            keyframes.append(frame)
        
        return keyframes
    
    def _create_happy_animation(self, skeleton: Dict) -> List[Dict]:
        """Create happy/excited animation keyframes"""
        # Jumping or bouncing motion
        keyframes = []
        
        for t in np.linspace(0, 2 * np.pi, 40):
            frame = {
                'time': t / (2 * np.pi),
                'bones': {}
            }
            
            # Bouncing motion
            bounce = abs(np.sin(t * 2)) * 0.1
            
            for bone in skeleton['bones']:
                frame['bones'][bone['id']] = {
                    'position': bone['position'] + [0, bounce, 0],
                    'rotation': [0, 0, 0, 1]
                }
            
            keyframes.append(frame)
        
        return keyframes
    
    def _save_rigged_model(self, rigged_model: Dict) -> str:
        """Save rigged model to file"""
        
        # Create temporary file
        with tempfile.NamedTemporaryFile(suffix='.glb', delete=False) as tmp:
            output_path = tmp.name
        
        # In production, this would export the rigged model with animations
        # For now, just save the mesh
        rigged_model['mesh'].export(output_path)
        
        return output_path
    
    def _save_mesh_without_rigging(self, mesh: Union[str, trimesh.Trimesh]) -> str:
        """Save mesh without rigging as fallback"""
        
        if isinstance(mesh, str):
            return mesh
        
        with tempfile.NamedTemporaryFile(suffix='.glb', delete=False) as tmp:
            output_path = tmp.name
        
        mesh.export(output_path)
        return output_path
    
    def to(self, device: str):
        """Move model to specified device (compatibility method)"""
        self.device = device