# SimpleJiggle - Vertex Jiggle Modifier for Blender
# Inspired by Marius Silaghi's Jiggle for 3ds Max
# Works directly on vertices without cloth physics or bones

bl_info = {
    "name": "SimpleJiggle",
    "author": "Your Name",
    "version": (1, 0, 0),
    "blender": (4, 0, 0),
    "location": "Properties > Modifiers",
    "description": "Adds jiggle-like animation to vertices based on object movement",
    "category": "Mesh",
}

import bpy
import numpy as np
from mathutils import Vector, Matrix
from bpy.props import (
    FloatProperty,
    StringProperty,
    BoolProperty,
    PointerProperty,
    IntProperty,
)
from bpy.types import (
    PropertyGroup,
    Operator,
    Operator,
    Panel,
    UIList,
)
from bpy.app.handlers import persistent


# ============================================================================
# JIGGLE DATA STORAGE
# ============================================================================

class JiggleLayerState:
    """Stores physics state for a single jiggle layer"""
    def __init__(self, vertex_count):
        self.jiggle_offset = np.zeros((vertex_count, 3), dtype=np.float64)
        self.jiggle_velocity = np.zeros((vertex_count, 3), dtype=np.float64)
        self.cached_weights = None
        self.cached_weights_col = None
        self.cached_vertex_group = None


class JiggleVertexData:
    """Stores per-object jiggle simulation data"""
    
    def __init__(self, vertex_count):
        self.vertex_count = vertex_count
        # Original REST positions
        self.rest_positions = np.zeros((vertex_count, 3), dtype=np.float64)
        # Previous frame's target world positions
        self.prev_world_positions = np.zeros((vertex_count, 3), dtype=np.float64)
        # Previous world matrix
        self.prev_world_matrix = None
        
        # Helper ones array
        self.ones = None
        
        # Name of the hidden source object
        self.source_obj_name = None
        self.initialized = False
        
        # Dictionary of layer states: {layer_name: JiggleLayerState}
        self.layers = {}

    def get_layer_state(self, layer_name):
        if layer_name not in self.layers:
            self.layers[layer_name] = JiggleLayerState(self.vertex_count)
        return self.layers[layer_name]


# Global storage for jiggle data per object
_jiggle_data = {}


def get_jiggle_data(obj):
    """Get or create jiggle data for an object"""
    global _jiggle_data
    obj_id = obj.name
    
    mesh = obj.data
    vertex_count = len(mesh.vertices)
    
    # Check if we need to create or recreate data
    if obj_id not in _jiggle_data:
        _jiggle_data[obj_id] = JiggleVertexData(vertex_count)
    elif _jiggle_data[obj_id].vertex_count != vertex_count:
        _jiggle_data[obj_id] = JiggleVertexData(vertex_count)
    
    return _jiggle_data[obj_id]


def clear_jiggle_data(obj):
    """Clear jiggle data for an object"""
    global _jiggle_data
    obj_id = obj.name
    if obj_id in _jiggle_data:
        del _jiggle_data[obj_id]


def set_armature_visibility(obj, visible):
    """Set visibility of all armature modifiers"""
    for mod in obj.modifiers:
        if mod.type == 'ARMATURE':
            mod.show_viewport = visible
            mod.show_render = visible


# ============================================================================
# JIGGLE SETTINGS (Per Object)
# ============================================================================

class SimpleJiggleSettings(PropertyGroup):
    """Settings for the SimpleJiggle modifier"""
    
    enabled: BoolProperty(
        name="Enable",
        description="Enable/Disable jiggle simulation",
        default=True,
    )
    
    jiggle_strength: FloatProperty(
        name="Jiggle",
        description="Overall strength of the jiggle effect (how much vertices react to movement)",
        default=1.0,
        min=0.0,
        max=2.0,
        soft_min=0.0,
        soft_max=1.0,
        subtype='FACTOR',
    )
    
    name: StringProperty(name="Name", default="Jiggle Layer")
    
    stiffness: FloatProperty(
        name="Stiffness",
        description="How quickly the jiggle springs back (higher = faster, snappier oscillation)",
        default=0.5,
        min=0.01,
        max=1.0,
        subtype='FACTOR',
    )
    
    damping: FloatProperty(
        name="Damping",
        description="How quickly the oscillation settles down (higher = fewer bounces)",
        default=0.3,
        min=0.0,
        max=1.0,
        subtype='FACTOR',
    )
    
    # NEW: Tension/Compression limits
    max_stretch: FloatProperty(
        name="Max Stretch",
        description="Maximum distance vertices can stretch from their target (in Blender units). Prevents extreme distortion",
        default=0.5,
        min=0.001,
        max=10.0,
        soft_min=0.01,
        soft_max=2.0,
        unit='LENGTH',
    )
    
    tension: FloatProperty(
        name="Tension",
        description="Spring tension - higher values create tighter, snappier springs with less wobble",
        default=0.5,
        min=0.0,
        max=1.0,
        subtype='FACTOR',
    )
    
    max_velocity: FloatProperty(
        name="Max Velocity",
        description="Maximum velocity of jiggle motion. Lower values create smoother, controlled movement",
        default=10.0,
        min=0.1,
        max=100.0,
        soft_min=1.0,
        soft_max=20.0,
    )
    
    vertex_group: StringProperty(
        name="Vertex Group",
        description="Vertex group to apply jiggle effect (weights control intensity)",
        default="",
    )
    
    # Internal tracking
    is_active: BoolProperty(
        name="Is Active",
        description="Whether the jiggle simulation is currently running",
        default=False,
    )


class SimpleJiggleProperties(PropertyGroup):
    """Container for multiple jiggle layers"""
    
    layers: bpy.props.CollectionProperty(type=SimpleJiggleSettings)
    active_index: IntProperty(name="Active Layer Index", default=0)
    
    # Global On/Off
    enabled: BoolProperty(name="Enable Jiggle", default=True)
    is_active: BoolProperty(name="Is Simulation Running", default=False)


# ============================================================================
# JIGGLE SIMULATION CORE
# ============================================================================

def get_vertex_weights(obj, vertex_group_name, layer_state):
    """
    Get vertex weights from a vertex group (with caching).
    Weights are cached in layer_state to avoid recalculating every frame.
    """
    # Check if we have cached weights for this vertex group
    if (layer_state.cached_weights is not None and 
        layer_state.cached_vertex_group == vertex_group_name):
        return layer_state.cached_weights, layer_state.cached_weights_col
    
    mesh = obj.data
    vertex_count = len(mesh.vertices)
    weights = np.zeros(vertex_count, dtype=np.float64)
    
    if not vertex_group_name:
        # No vertex group specified, use weight of 1.0 for all vertices
        weights.fill(1.0)
    else:
        # Find the vertex group
        vg = obj.vertex_groups.get(vertex_group_name)
        if vg is not None:
            vg_index = vg.index
            # Get weights for each vertex
            for i, vert in enumerate(mesh.vertices):
                for group in vert.groups:
                    if group.group == vg_index:
                        weights[i] = group.weight
                        break
    
    # Cache the weights
    weights_col = weights[:, np.newaxis]  # Shape: (N, 1) for broadcasting
    layer_state.cached_weights = weights
    layer_state.cached_weights_col = weights_col
    layer_state.cached_vertex_group = vertex_group_name
    
    return weights, weights_col


def simulate_jiggle_post(obj, delta_time, depsgraph):
    """
    Run jiggle simulation (POST-FRAME phase).
    
    For SKINNED meshes: PRE handler already restored REST, so evaluated mesh
    now has clean armature output.
    
    For UNSKINNED meshes: Use stored REST positions directly.
    """
    settings = obj.simple_jiggle
    
    if not settings.enabled or not settings.is_active:
        return
    
    if len(settings.layers) == 0:
        return
        
    mesh = obj.data
    jiggle_data = get_jiggle_data(obj)
    vertex_count = len(mesh.vertices)
    
    # Check if there's an armature modifier
    has_armature = any(mod.type == 'ARMATURE' for mod in obj.modifiers)
    
    # Get world matrix
    world_matrix_np = np.array(obj.matrix_world, dtype=np.float64)
    world_matrix_inv_np = np.linalg.inv(world_matrix_np)
    ones = np.ones((vertex_count, 1), dtype=np.float64)
    
    # =========================================================================
    # INITIALIZATION: Store original REST positions (only once!)
    # =========================================================================
    
    
    # Get source object
    source_obj = bpy.data.objects.get(jiggle_data.source_obj_name)
    if not source_obj:
        # Source missing or deleted, stop simulation
        settings.is_active = False
        return

    if not jiggle_data.initialized:
        # Read the CURRENT mesh positions as REST
        # IMPORTANT: This should be called when mesh is in clean state
        mesh.vertices.foreach_get('co', jiggle_data.rest_positions.ravel())
        
        # Calculate initial world positions for tracking
        rest_h = np.hstack([jiggle_data.rest_positions, ones])
        rest_world = (rest_h @ world_matrix_np.T)[:, :3]
        
        jiggle_data.prev_world_positions = rest_world.copy()
        jiggle_data.prev_world_matrix = world_matrix_np.copy()
        jiggle_data.ones = ones
        jiggle_data.initialized = True
        return
    
    # =========================================================================
    # Calculate TARGET positions from HIDDEN SOURCE
    # =========================================================================
    
    # We always read from the source object, which has the correct animation/armature state
    source_eval = source_obj.evaluated_get(depsgraph)
    source_mesh = source_eval.data
    
    target_local = np.zeros((vertex_count, 3), dtype=np.float64)
    # Ensure vertex counts match (safety check)
    if len(source_mesh.vertices) != vertex_count:
        settings.is_active = False
        return
        
    source_mesh.vertices.foreach_get('co', target_local.ravel())
    
    # Helper: Get source world matrix (should be same as obj but safer to read)
    source_matrix_np = np.array(source_obj.matrix_world, dtype=np.float64)
    
    # Convert target to world space
    target_h = np.hstack([target_local, ones])
    target_world = (target_h @ source_matrix_np.T)[:, :3]
    
    # Calculate movement (how much target positions moved since last frame)
    movement_delta = target_world - jiggle_data.prev_world_positions
    dt = min(delta_time, 1.0 / 24.0)
    
    # Store accumulated offset for all layers
    total_weighted_offset = np.zeros((vertex_count, 3), dtype=np.float64)
    
    # Iterate through all enabled layers
    for i, layer_settings in enumerate(settings.layers):
        if not layer_settings.enabled:
            continue
            
        # Get physics state for this layer
        layer_state = jiggle_data.get_layer_state(str(i)) # Use index as stable ID during runtime
        
        # Get weights
        weights, weights_col = get_vertex_weights(obj, layer_settings.vertex_group, layer_state)
        
        # --- PHYSICS SIMULATION FOR THIS LAYER ---
        
        base_stiffness = 100.0 + (layer_settings.stiffness * 400.0)
        tension_mult = 1.0 + (layer_settings.tension * 4.0)
        stiffness_k = base_stiffness * tension_mult
        
        base_damping = 5.0 + (layer_settings.damping * 25.0)
        damping_c = base_damping * (1.0 + layer_settings.tension * 0.5)
        
        jiggle_mult = layer_settings.jiggle_strength
        max_stretch = layer_settings.max_stretch
        max_velocity = layer_settings.max_velocity
        
        # Apply impulse
        impulse_strength = jiggle_mult * 30.0 * (1.0 + layer_settings.tension * 0.5)
        layer_state.jiggle_velocity -= movement_delta * impulse_strength
        
        # Velocity limiting
        velocity_magnitude = np.linalg.norm(layer_state.jiggle_velocity, axis=1, keepdims=True)
        velocity_scale = np.where(
            velocity_magnitude > max_velocity,
            max_velocity / (velocity_magnitude + 1e-10),
            1.0
        )
        layer_state.jiggle_velocity *= velocity_scale
        
        # Spring physics
        acceleration = (-stiffness_k * layer_state.jiggle_offset) + (-damping_c * layer_state.jiggle_velocity)
        layer_state.jiggle_velocity += acceleration * dt
        layer_state.jiggle_offset += layer_state.jiggle_velocity * dt
        
        # Stretch limiting
        offset_magnitude = np.linalg.norm(layer_state.jiggle_offset, axis=1, keepdims=True)
        stretch_exceeded = offset_magnitude > max_stretch
        offset_scale = np.where(
            stretch_exceeded,
            max_stretch / (offset_magnitude + 1e-10),
            1.0
        )
        layer_state.jiggle_offset *= offset_scale
        
        if np.any(stretch_exceeded):
            layer_state.jiggle_velocity *= np.where(stretch_exceeded, 0.5, 1.0)
            
        # Accumulate to total offset (weighted)
        total_weighted_offset += layer_state.jiggle_offset * weights_col


    # Store global state for next frame
    jiggle_data.prev_world_positions = target_world.copy()
    jiggle_data.prev_world_matrix = world_matrix_np.copy()
    
    # =========================================================================
    # FINAL POSITIONS: target + accumulated weighted offsets
    # =========================================================================
    
    final_world = target_world + total_weighted_offset
    
    # Convert back to local space
    final_world_h = np.hstack([final_world, ones])
    final_local = (final_world_h @ world_matrix_inv_np.T)[:, :3]
    
    # Write to mesh
    mesh.vertices.foreach_set('co', final_local.ravel().astype(np.float32))
    mesh.update()
    
    # Note: We do NOT need to toggle armature visibility anymore because
    # the active object has its armature disabled during simulation.


# ============================================================================
# OPERATORS
# ============================================================================

class SIMPLEJIGGLE_OT_start(Operator):
    """Start jiggle simulation"""
    bl_idname = "simplejiggle.start"
    bl_label = "Start Jiggle"
    bl_options = {'REGISTER', 'UNDO'}
    
    @classmethod
    def poll(cls, context):
        obj = context.active_object
        return obj is not None and obj.type == 'MESH'
    
    def execute(self, context):
        obj = context.active_object
        settings = obj.simple_jiggle
        
        # Clear existing data and reinitialize
        clear_jiggle_data(obj)
        jiggle_data = get_jiggle_data(obj)
        
        # Create hidden source object
        source_name = f"{obj.name}_JiggleSource"
        
        # Cleanup existing if any (crashed session?)
        old_source = bpy.data.objects.get(source_name)
        if old_source:
            bpy.data.objects.remove(old_source, do_unlink=True)
            
        # Create duplicate
        source_obj = obj.copy()
        source_obj.name = source_name
        source_obj.data = obj.data.copy() # Deep copy data to separate it
        context.collection.objects.link(source_obj)
        
        # Setup source object (Visible to depsgraph, hidden from view)
        # IMPORTANT: Object MUST be visible in viewport to be evaluated by depsgraph in newer Blender versions!
        # We use 'BOUNDS' display type to make it unobtrusive.
        source_obj.hide_viewport = False 
        source_obj.display_type = 'BOUNDS'
        source_obj.hide_render = True
        source_obj.hide_select = True
        
        # Store reference
        jiggle_data.source_obj_name = source_name
        
        # Disable Armature on ACTIVE object (Result), ensure enabled on SOURCE
        set_armature_visibility(obj, False)
        set_armature_visibility(source_obj, True)
        
        settings.is_active = True
        
        # If no layers, allow starting but maybe warn?
        if len(settings.layers) == 0:
            self.report({'WARNING'}, "Jiggle started but no layers added")
        
        self.report({'INFO'}, "Jiggle simulation started")
        return {'FINISHED'}


class SIMPLEJIGGLE_OT_stop(Operator):
    """Stop jiggle simulation"""
    bl_idname = "simplejiggle.stop"
    bl_label = "Stop Jiggle"
    bl_options = {'REGISTER', 'UNDO'}
    
    @classmethod
    def poll(cls, context):
        obj = context.active_object
        return obj is not None and obj.type == 'MESH'
    
    def execute(self, context):
        obj = context.active_object
        settings = obj.simple_jiggle
        settings.is_active = False
        
        # Restore armature visibility
        set_armature_visibility(obj, True)
        
        # Restore REST positions (Fix the mesh)
        jiggle_data = get_jiggle_data(obj)
        if jiggle_data.initialized:
            mesh = obj.data
            try:
                mesh.vertices.foreach_set('co', jiggle_data.rest_positions.ravel().astype(np.float32))
                mesh.update()
            except:
                pass
        
        # Cleanup source object
        if hasattr(jiggle_data, 'source_obj_name') and jiggle_data.source_obj_name:
            source_obj = bpy.data.objects.get(jiggle_data.source_obj_name)
            if source_obj:
                bpy.data.objects.remove(source_obj, do_unlink=True)
            jiggle_data.source_obj_name = None

        clear_jiggle_data(obj)
        
        self.report({'INFO'}, "Jiggle simulation stopped")
        return {'FINISHED'}


class SIMPLEJIGGLE_OT_reset(Operator):
    """Reset jiggle simulation and restore original mesh"""
    bl_idname = "simplejiggle.reset"
    bl_label = "Reset Jiggle"
    bl_options = {'REGISTER', 'UNDO'}
    
    @classmethod
    def poll(cls, context):
        obj = context.active_object
        return obj is not None and obj.type == 'MESH'
    
    def execute(self, context):
        obj = context.active_object
        settings = obj.simple_jiggle
        settings.is_active = False
        clear_jiggle_data(obj)
        
        # Restore armature visibility
        set_armature_visibility(obj, True)
        
        self.report({'INFO'}, "Jiggle simulation reset")
        return {'FINISHED'}


class SIMPLEJIGGLE_OT_add_layer(Operator):
    """Add a new jiggle layer"""
    bl_idname = "simplejiggle.add_layer"
    bl_label = "Add Layer"
    bl_options = {'REGISTER', 'UNDO'}
    
    def execute(self, context):
        obj = context.active_object
        settings = obj.simple_jiggle
        
        item = settings.layers.add()
        item.name = f"Layer {len(settings.layers)}"
        settings.active_index = len(settings.layers) - 1
        
        return {'FINISHED'}


class SIMPLEJIGGLE_OT_remove_layer(Operator):
    """Remove active jiggle layer"""
    bl_idname = "simplejiggle.remove_layer"
    bl_label = "Remove Layer"
    bl_options = {'REGISTER', 'UNDO'}
    
    def execute(self, context):
        obj = context.active_object
        settings = obj.simple_jiggle
        
        if settings.layers:
            settings.layers.remove(settings.active_index)
            settings.active_index = min(max(0, settings.active_index - 1), len(settings.layers) - 1)
            
        return {'FINISHED'}


class SIMPLEJIGGLE_OT_bake(Operator):
    """Bake jiggle animation to shape keys"""
    bl_idname = "simplejiggle.bake"
    bl_label = "Bake to Shape Keys"
    bl_options = {'REGISTER', 'UNDO'}
    
    @classmethod
    def poll(cls, context):
        obj = context.active_object
        return obj is not None and obj.type == 'MESH' and obj.simple_jiggle.is_active
    
    def execute(self, context):
        obj = context.active_object
        scene = context.scene
        
        # Store current frame
        original_frame = scene.frame_current
        
        # Ensure we have a basis shape key
        if obj.data.shape_keys is None:
            obj.shape_key_add(name="Basis", from_mix=False)
        
        # Bake each frame
        frame_start = scene.frame_start
        frame_end = scene.frame_end
        
        for frame in range(frame_start, frame_end + 1):
            scene.frame_set(frame)
            
            # Create shape key for this frame
            sk = obj.shape_key_add(name=f"Jiggle_{frame}", from_mix=False)
            
            # Copy current vertex positions to shape key
            mesh = obj.data
            for i, vert in enumerate(mesh.vertices):
                sk.data[i].co = vert.co.copy()
            
            # Insert keyframe for value
            sk.value = 0.0
            sk.keyframe_insert(data_path="value", frame=frame - 1)
            sk.value = 1.0
            sk.keyframe_insert(data_path="value", frame=frame)
            sk.value = 0.0
            sk.keyframe_insert(data_path="value", frame=frame + 1)
        
        # Restore original frame
        scene.frame_set(original_frame)
        
        # Restore original frame
        scene.frame_set(original_frame)
        
        # Stop simulation properly to cleanup
        bpy.ops.simplejiggle.stop()
        
        self.report({'INFO'}, f"Baked jiggle to {frame_end - frame_start + 1} shape keys")
        return {'FINISHED'}


# ============================================================================
# FRAME CHANGE HANDLERS (Two-phase for skinned mesh support)
# ============================================================================

_last_frame = -1
_last_time = 0.0

@persistent
def jiggle_frame_change_pre(scene):
    """
    PRE-FRAME HANDLER: Restore REST positions for skinned meshes.
    
    This runs BEFORE Blender evaluates the armature, so the armature
    will see clean REST positions and produce clean output.
    """
    global _jiggle_data
    
    # Handler removed: We no longer need to reset mesh in PRE 
    # because we use a separate source object for clean input.
    pass


@persistent
def jiggle_frame_change_post(scene, depsgraph):
    """
    POST-FRAME HANDLER: Read armature output and apply jiggle.
    
    This runs AFTER Blender evaluates the armature, so we get clean
    armature-deformed positions to apply jiggle on top of.
    """
    global _last_frame, _last_time
    
    current_frame = scene.frame_current
    fps = scene.render.fps / scene.render.fps_base
    
    # Calculate delta time
    if current_frame != _last_frame:
        delta_time = abs(current_frame - _last_frame) / fps
        _last_frame = current_frame
    else:
        delta_time = 1.0 / fps
    
    # Process all objects with jiggle enabled
    for obj in scene.objects:
        if obj.type != 'MESH':
            continue
        
        settings = obj.simple_jiggle
        if settings.is_active and settings.enabled:
            try:
                simulate_jiggle_post(obj, delta_time, depsgraph)
            except Exception as e:
                print(f"SimpleJiggle POST Error on {obj.name}: {e}")


# ============================================================================
# UI PANEL
# ============================================================================

class SIMPLEJIGGLE_UL_layer_list(UIList):
    def draw_item(self, context, layout, data, item, icon, active_data, active_propname, index):
        if self.layout_type in {'DEFAULT', 'COMPACT'}:
            layout.prop(item, "enabled", text="")
            layout.prop(item, "name", text="", emboss=False)
            layout.prop_search(item, "vertex_group", data, "vertex_groups", text="", icon='GROUP_VERTEX')
        elif self.layout_type == 'GRID':
            layout.alignment = 'CENTER'
            layout.label(text="", icon='FORCE_HARMONIC')

class SIMPLEJIGGLE_PT_main(Panel):
    """Main panel for SimpleJiggle settings"""
    bl_label = "SimpleJiggle"
    bl_idname = "SIMPLEJIGGLE_PT_main"
    bl_space_type = 'PROPERTIES'
    bl_region_type = 'WINDOW'
    bl_context = "physics"
    bl_options = {'DEFAULT_CLOSED'}
    
    @classmethod
    def poll(cls, context):
        obj = context.active_object
        return obj is not None and obj.type == 'MESH'
    
    def draw_header(self, context):
        obj = context.active_object
        settings = obj.simple_jiggle
        self.layout.prop(settings, "enabled", text="")
    
    def draw(self, context):
        layout = self.layout
        layout.use_property_split = True
        layout.use_property_decorate = False
        
        obj = context.active_object
        settings = obj.simple_jiggle
        
        # Enable/Disable based on main toggle
        layout.enabled = settings.enabled
        
        # Status indicator
        box = layout.box()
        row = box.row()
        if settings.is_active:
            row.label(text="Status: Active", icon='PLAY')
        else:
            row.label(text="Status: Inactive", icon='PAUSE')
        
        # Control buttons
        row = box.row(align=True)
        if not settings.is_active:
            row.operator("simplejiggle.start", text="Start", icon='PLAY')
        else:
            row.operator("simplejiggle.stop", text="Stop", icon='PAUSE')
        row.operator("simplejiggle.reset", text="Reset", icon='FILE_REFRESH')
        
        layout.separator()
        
        # Layers List
        row = layout.row()
        row.template_list("SIMPLEJIGGLE_UL_layer_list", "", settings, "layers", settings, "active_index")
        
        col = row.column(align=True)
        col.operator("simplejiggle.add_layer", icon='ADD', text="")
        col.operator("simplejiggle.remove_layer", icon='REMOVE', text="")
        
        # Parameters for active layer
        if settings.layers and settings.active_index >= 0 and settings.active_index < len(settings.layers):
            active_layer = settings.layers[settings.active_index]
            
            box = layout.box()
            box.label(text=f"Settings: {active_layer.name}", icon='PREFERENCES')
            
            col = box.column(align=True)
            col.prop(active_layer, "jiggle_strength", text="Jiggle Amount")
            col.prop(active_layer, "stiffness")
            col.prop(active_layer, "damping")
            col.prop(active_layer, "tension")
            
            col = box.column(align=True)
            col.prop(active_layer, "max_stretch")
            col.prop(active_layer, "max_velocity")
            
            box.prop_search(active_layer, "vertex_group", obj, "vertex_groups", text="Vertex Group")
            
        
        # Bake section
        layout.separator()
        box = layout.box()
        box.label(text="Baking", icon='RENDER_ANIMATION')
        box.operator("simplejiggle.bake", text="Bake to Shape Keys", icon='SHAPEKEY_DATA')


# ============================================================================
# REGISTER / UNREGISTER
# ============================================================================

classes = (
    SimpleJiggleSettings,
    SimpleJiggleProperties,
    SIMPLEJIGGLE_OT_start,
    SIMPLEJIGGLE_OT_stop,
    SIMPLEJIGGLE_OT_reset,
    SIMPLEJIGGLE_OT_add_layer,
    SIMPLEJIGGLE_OT_remove_layer,
    SIMPLEJIGGLE_OT_bake,
    SIMPLEJIGGLE_UL_layer_list,
    SIMPLEJIGGLE_PT_main,
)


def register():
    for cls in classes:
        bpy.utils.register_class(cls)
    
    # Add settings to Object type
    bpy.types.Object.simple_jiggle = PointerProperty(type=SimpleJiggleProperties)
    
    # Add frame change handlers (PRE restores REST, POST applies jiggle)
    bpy.app.handlers.frame_change_post.append(jiggle_frame_change_post)
    
    print("SimpleJiggle registered successfully")


def unregister():
    # Remove frame change handlers
    if jiggle_frame_change_post in bpy.app.handlers.frame_change_post:
        bpy.app.handlers.frame_change_post.remove(jiggle_frame_change_post)
    
    # Clear global data
    global _jiggle_data
    _jiggle_data.clear()
    
    # Remove settings from Object type
    del bpy.types.Object.simple_jiggle
    
    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)
    
    print("SimpleJiggle unregistered successfully")


if __name__ == "__main__":
    register()
