blender_splat_export.py

bl_info = {
    "name": "Gaussian Splat Exporter",
    "author": "PLAN8",
    "version": (0, 3, 0),
    "blender": (3, 0, 0),
    "location": "View3D > Sidebar > Gaussian Splat | File > Export > Gaussian Splat (.ply/.sog)",
    "description": "Export mesh geometry to Gaussian Splat format with direct PLY export or via splat-transform",
    "category": "Import-Export",
}

import bpy
import bmesh
import json
import os
import subprocess
from bpy.props import StringProperty, FloatProperty, IntProperty, BoolProperty, EnumProperty, PointerProperty
from bpy_extras.io_utils import ExportHelper
from mathutils import Vector, Matrix
import math


class GaussianSplatSettings(bpy.types.PropertyGroup):
    """Settings for Gaussian Splat export"""
    
    export_format: EnumProperty(
        name="Export Format",
        description="Choose the output file format",
        items=(
            ('PLY', ".ply", "Export as PLY format"),
            ('SOG', ".sog", "Export as SOG format"),
        ),
        default='PLY',
    )
    
    use_splat_transform: BoolProperty(
        name="Use splat-transform",
        description="Use external splat-transform tool (requires installation). If disabled, writes PLY directly",
        default=False
    )

    splat_transform_path: StringProperty(
        name="splat-transform Command",
        description="Command to run splat-transform (e.g., 'splat-transform' if globally installed, or full path)",
        default="splat-transform",
    )

    overwrite_output: BoolProperty(
        name="Overwrite Existing File",
        description="Overwrite output file if it exists",
        default=True
    )

    keep_mjs_file: BoolProperty(
        name="Export .mjs File",
        description="Save the .mjs generator file in the export directory (only used with splat-transform)",
        default=False
    )

    use_frame_number: BoolProperty(
        name="Use Frame Number in Filename",
        description="Append the current frame number to the exported filename",
        default=False
    )

    batch_export_animation: BoolProperty(
        name="Batch Export Animation",
        description="Export all frames from timeline start to end",
        default=False
    )

    sample_density: FloatProperty(
        name="Sample Density",
        description="Number of splats per square unit (only used in Surface Sampling mode)",
        default=100.0,
        min=1.0,
        max=10000.0
    )

    sampling_mode: EnumProperty(
        name="Sampling Mode",
        description="How to generate splats from the mesh",
        items=(
            ('VERTICES', "Use Vertices", "Create one splat per vertex"),
            ('SURFACE', "Surface Sampling", "Sample splats across surface area"),
        ),
        default='VERTICES',
    )

    splat_scale: FloatProperty(
        name="Global Scale Multiplier",
        description="Global scale multiplier applied to all splats (multiplied with auto-calculated size based on vertex proximity)",
        default=1.0,
        min=0.001,
        max=100.0
    )
    
    use_auto_scale: BoolProperty(
        name="Auto Scale from Vertex Proximity",
        description="Automatically scale splats based on distance to nearest vertices",
        default=True
    )

    auto_scale_neighbors: IntProperty(
        name="Auto Scale Neighbor Count",
        description="Number of nearest neighbors to average for auto-scaling (higher = smoother)",
        default=5,
        min=1,
        max=20
    )

    splat_opacity: FloatProperty(
        name="Global Opacity Multiplier",
        description="Global opacity multiplier applied to all splats (multiplied with color attribute alpha)",
        default=1.0,
        min=0.0,
        max=1.0
    )

    use_vertex_colors: BoolProperty(
        name="Use Vertex Colors",
        description="Use vertex colors if available",
        default=True
    )

    use_normals: BoolProperty(
        name="Orient to Normals",
        description="Orient splats perpendicular to surface normals",
        default=True
    )

    axis_forward: EnumProperty(
        name="Forward",
        items=(
            ('X', "X Forward", ""),
            ('Y', "Y Forward", ""),
            ('Z', "Z Forward", ""),
            ('-X', "-X Forward", ""),
            ('-Y', "-Y Forward", ""),
            ('-Z', "-Z Forward", ""),
        ),
        default='Z',
    )

    axis_up: EnumProperty(
        name="Up",
        items=(
            ('X', "X Up", ""),
            ('Y', "Y Up", ""),
            ('Z', "Z Up", ""),
            ('-X', "-X Up", ""),
            ('-Y', "-Y Up", ""),
            ('-Z', "-Z Up", ""),
        ),
        default='-Y',
    )

    export_path: StringProperty(
        name="Export Path",
        description="Path for the exported file",
        default="",
        subtype='FILE_PATH'
    )


class GAUSSIANSPLAT_PT_MainPanel(bpy.types.Panel):
    """Main panel for Gaussian Splat export in 3D view sidebar"""
    bl_label = "Gaussian Splat Export"
    bl_idname = "GAUSSIANSPLAT_PT_main_panel"
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category = 'Gaussian Splat'

    def draw(self, context):
        layout = self.layout
        settings = context.scene.gaussian_splat_settings
        
        # Export path
        box = layout.box()
        box.label(text="Export Settings:", icon='EXPORT')
        box.prop(settings, "export_path")
        box.prop(settings, "export_format")
        box.prop(settings, "use_splat_transform")
        if settings.use_splat_transform:
            box.prop(settings, "splat_transform_path")
            box.prop(settings, "keep_mjs_file")
        box.prop(settings, "overwrite_output")
        
        # Animation export options
        box = layout.box()
        box.label(text="Animation Export:", icon='TIME')
        box.prop(settings, "use_frame_number")
        box.prop(settings, "batch_export_animation")
        if settings.batch_export_animation:
            row = box.row()
            row.label(text=f"Will export frames {context.scene.frame_start} to {context.scene.frame_end}")
        
        # Sampling options
        box = layout.box()
        box.label(text="Sampling Options:", icon='MESH_DATA')
        box.prop(settings, "sampling_mode")
        if settings.sampling_mode == 'SURFACE':
            box.prop(settings, "sample_density")
        
        # Splat properties
        box = layout.box()
        box.label(text="Splat Properties:", icon='PARTICLE_DATA')
        box.prop(settings, "use_auto_scale")
        if settings.use_auto_scale:
            box.prop(settings, "auto_scale_neighbors")
        box.prop(settings, "splat_scale")
        box.prop(settings, "splat_opacity")
        
        # Color options
        box = layout.box()
        box.label(text="Color Options:", icon='COLOR')
        box.prop(settings, "use_vertex_colors")
        box.prop(settings, "use_normals")
        
        # Axis conversion
        box = layout.box()
        box.label(text="Axis Conversion:", icon='ORIENTATION_GIMBAL')
        box.prop(settings, "axis_forward")
        box.prop(settings, "axis_up")
        
        # Export button
        layout.separator()
        row = layout.row(align=True)
        row.scale_y = 1.5
        row.operator("export_scene.gaussian_splat_direct", text="Export Gaussian Splat", icon='EXPORT')


class GAUSSIANSPLAT_OT_DirectExport(bpy.types.Operator):
    """Export Gaussian Splat directly from the side panel"""
    bl_idname = "export_scene.gaussian_splat_direct"
    bl_label = "Export Gaussian Splat"
    bl_options = {'REGISTER'}

    def execute(self, context):
        settings = context.scene.gaussian_splat_settings
        
        # Validate export path
        if not settings.export_path:
            self.report({'ERROR'}, "Please specify an export path")
            return {'CANCELLED'}
        
        # Get selected objects
        selected_objects = [obj for obj in context.selected_objects if obj.type == 'MESH']

        if not selected_objects:
            self.report({'ERROR'}, "No mesh objects selected")
            return {'CANCELLED'}

        # Check if batch export is enabled
        if settings.batch_export_animation:
            return self.batch_export_frames(context, settings, selected_objects)
        else:
            return self.export_single_frame(context, settings, selected_objects, context.scene.frame_current)

    def export_single_frame(self, context, settings, selected_objects, frame_number):
        """Export a single frame"""
        # Set the frame
        context.scene.frame_set(frame_number)
        
        # Get file extension based on format
        file_ext = '.ply' if settings.export_format == 'PLY' else '.sog'
        
        # Build filepath with optional frame number
        base_path = settings.export_path
        # Remove existing extension if present
        for ext in ['.ply', '.sog']:
            if base_path.lower().endswith(ext):
                base_path = base_path[:-len(ext)]
                break
        
        if settings.use_frame_number:
            # Insert frame number before extension
            filepath = f"{base_path}_{frame_number:04d}{file_ext}"
        else:
            filepath = base_path + file_ext
        
        # Derive output directory
        export_dir = os.path.dirname(filepath)
        if not export_dir:
            export_dir = bpy.path.abspath("//")
            filepath = os.path.join(export_dir, os.path.basename(filepath))

        try:
            # Get axis conversion matrix
            axis_matrix = get_axis_conversion_matrix(settings)

            # Sample points from meshes
            all_samples = []
            for obj in selected_objects:
                if settings.sampling_mode == 'VERTICES':
                    samples = sample_vertices(obj, context, axis_matrix, settings)
                else:
                    samples = sample_mesh(obj, context, axis_matrix, settings)
                all_samples.extend(samples)

            self.report({'INFO'}, f"Frame {frame_number}: Sampled {len(all_samples)} points from {len(selected_objects)} object(s)")

            # Choose export method
            if settings.use_splat_transform:
                # Use splat-transform method
                base_name = os.path.splitext(os.path.basename(filepath))[0]
                generator_path = os.path.join(export_dir, f"{base_name}.mjs")
                
                create_mesh_generator(generator_path, all_samples)

                # Build splat-transform command
                cmd = [settings.splat_transform_path]
                if settings.overwrite_output:
                    cmd.append('-w')
                cmd.extend([generator_path, filepath])
                
                # Run splat-transform
                result = subprocess.run(
                    cmd,
                    capture_output=True,
                    text=True,
                    shell=True
                )

                if result.returncode != 0:
                    self.report({'ERROR'}, f"splat-transform failed: {result.stderr}")
                    return {'CANCELLED'}

                # Delete .mjs file if not keeping it
                if not settings.keep_mjs_file and os.path.exists(generator_path):
                    try:
                        os.remove(generator_path)
                    except Exception as e:
                        self.report({'WARNING'}, f"Could not delete .mjs file: {str(e)}")
            else:
                # Write PLY directly
                if settings.export_format == 'PLY':
                    write_ply_direct(filepath, all_samples, settings.overwrite_output)
                else:
                    self.report({'ERROR'}, "Direct export only supports PLY format. Use splat-transform for SOG format.")
                    return {'CANCELLED'}

            self.report({'INFO'}, f"Successfully exported to {filepath}")
            return {'FINISHED'}

        except Exception as e:
            import traceback
            self.report({'ERROR'}, f"Export failed: {str(e)}")
            print(traceback.format_exc())  # Print full traceback to console
            return {'CANCELLED'}

    def batch_export_frames(self, context, settings, selected_objects):
        """Export all frames in the timeline range"""
        start_frame = context.scene.frame_start
        end_frame = context.scene.frame_end
        original_frame = context.scene.frame_current
        
        total_frames = end_frame - start_frame + 1
        self.report({'INFO'}, f"Starting batch export of {total_frames} frames...")
        
        success_count = 0
        fail_count = 0
        
        for frame in range(start_frame, end_frame + 1):
            result = self.export_single_frame(context, settings, selected_objects, frame)
            if result == {'FINISHED'}:
                success_count += 1
            else:
                fail_count += 1
        
        # Restore original frame
        context.scene.frame_set(original_frame)
        
        if fail_count == 0:
            self.report({'INFO'}, f"Batch export complete! Successfully exported {success_count} frames.")
            return {'FINISHED'}
        else:
            self.report({'WARNING'}, f"Batch export finished with errors. Success: {success_count}, Failed: {fail_count}")
            return {'FINISHED'}


class GaussianSplatExporter(bpy.types.Operator, ExportHelper):
    """Export mesh to Gaussian Splat format"""
    bl_idname = "export_scene.gaussian_splat"
    bl_label = "Export Gaussian Splat"

    filename_ext = ".ply"

    filter_glob: StringProperty(
        default="*.ply;*.sog",
        options={'HIDDEN'},
    )

    export_format: EnumProperty(
        name="Export Format",
        description="Choose the output file format",
        items=(
            ('PLY', ".ply", "Export as PLY format"),
            ('SOG', ".sog", "Export as SOG format"),
        ),
        default='PLY',
    )

    use_splat_transform: BoolProperty(
        name="Use splat-transform",
        description="Use external splat-transform tool (requires installation). If disabled, writes PLY directly",
        default=False
    )

    splat_transform_path: StringProperty(
        name="splat-transform Command",
        description="Command to run splat-transform (e.g., 'splat-transform' if globally installed, or full path)",
        default="splat-transform",
    )

    overwrite_output: BoolProperty(
        name="Overwrite Existing File",
        description="Overwrite output file if it exists",
        default=True
    )

    keep_mjs_file: BoolProperty(
        name="Export .mjs File",
        description="Save the .mjs generator file in the export directory (only used with splat-transform)",
        default=False
    )

    use_frame_number: BoolProperty(
        name="Use Frame Number in Filename",
        description="Append the current frame number to the exported filename",
        default=False
    )

    batch_export_animation: BoolProperty(
        name="Batch Export Animation",
        description="Export all frames from timeline start to end",
        default=False
    )

    sample_density: FloatProperty(
        name="Sample Density",
        description="Number of splats per square unit (only used in Surface Sampling mode)",
        default=100.0,
        min=1.0,
        max=10000.0
    )

    sampling_mode: EnumProperty(
        name="Sampling Mode",
        description="How to generate splats from the mesh",
        items=(
            ('VERTICES', "Use Vertices", "Create one splat per vertex"),
            ('SURFACE', "Surface Sampling", "Sample splats across surface area"),
        ),
        default='VERTICES',
    )

    splat_scale: FloatProperty(
        name="Global Scale Multiplier",
        description="Global scale multiplier applied to all splats (multiplied with auto-calculated size based on vertex proximity)",
        default=1.0,
        min=0.001,
        max=1000.0
    )
    
    use_auto_scale: BoolProperty(
        name="Auto Scale from Vertex Proximity",
        description="Automatically scale splats based on distance to nearest vertices",
        default=True
    )

    auto_scale_neighbors: IntProperty(
        name="Auto Scale Neighbor Count",
        description="Number of nearest neighbors to average for auto-scaling (higher = smoother)",
        default=5,
        min=1,
        max=20
    )

    splat_opacity: FloatProperty(
        name="Global Opacity Multiplier",
        description="Global opacity multiplier applied to all splats (multiplied with color attribute alpha)",
        default=1.0,
        min=0.0,
        max=1.0
    )

    use_vertex_colors: BoolProperty(
        name="Use Vertex Colors",
        description="Use vertex colors if available",
        default=True
    )

    use_normals: BoolProperty(
        name="Orient to Normals",
        description="Orient splats perpendicular to surface normals",
        default=True
    )

    axis_forward: EnumProperty(
        name="Forward",
        items=(
            ('X', "X Forward", ""),
            ('Y', "Y Forward", ""),
            ('Z', "Z Forward", ""),
            ('-X', "-X Forward", ""),
            ('-Y', "-Y Forward", ""),
            ('-Z', "-Z Forward", ""),
        ),
        default='-Z',
    )

    axis_up: EnumProperty(
        name="Up",
        items=(
            ('X', "X Up", ""),
            ('Y', "Y Up", ""),
            ('Z', "Z Up", ""),
            ('-X', "-X Up", ""),
            ('-Y', "-Y Up", ""),
            ('-Z', "-Z Up", ""),
        ),
        default='Y',
    )

    def draw(self, context):
        """Draw the export options in the file browser"""
        layout = self.layout
        
        layout.prop(self, "use_splat_transform")
        if self.use_splat_transform:
            layout.prop(self, "splat_transform_path")
            layout.prop(self, "keep_mjs_file")
        layout.prop(self, "export_format")
        layout.prop(self, "overwrite_output")
        
        layout.separator()
        layout.label(text="Animation Export:")
        layout.prop(self, "use_frame_number")
        layout.prop(self, "batch_export_animation")
        if self.batch_export_animation:
            layout.label(text=f"Will export frames {context.scene.frame_start} to {context.scene.frame_end}")
        
        layout.separator()
        layout.label(text="Sampling Options:")
        layout.prop(self, "sampling_mode")
        if self.sampling_mode == 'SURFACE':
            layout.prop(self, "sample_density")
        
        layout.separator()
        layout.label(text="Splat Properties:")
        layout.prop(self, "use_auto_scale")
        if self.use_auto_scale:
            layout.prop(self, "auto_scale_neighbors")
        layout.prop(self, "splat_scale")
        layout.prop(self, "splat_opacity")
        
        layout.separator()
        layout.label(text="Color Options:")
        layout.prop(self, "use_vertex_colors")
        layout.prop(self, "use_normals")
        
        layout.separator()
        layout.label(text="Axis Conversion:")
        layout.prop(self, "axis_forward")
        layout.prop(self, "axis_up")

    def execute(self, context):
        selected_objects = [obj for obj in context.selected_objects if obj.type == 'MESH']

        if not selected_objects:
            self.report({'ERROR'}, "No mesh objects selected")
            return {'CANCELLED'}

        # Check if batch export is enabled
        if self.batch_export_animation:
            return self.batch_export_frames(context, selected_objects)
        else:
            return self.export_single_frame(context, selected_objects, context.scene.frame_current)

    def export_single_frame(self, context, selected_objects, frame_number):
        """Export a single frame"""
        # Set the frame
        context.scene.frame_set(frame_number)
        
        # Get file extension based on format
        file_ext = '.ply' if self.export_format == 'PLY' else '.sog'
        
        # Build filepath with optional frame number and correct extension
        base_path = self.filepath
        # Remove existing extension if present
        for ext in ['.ply', '.sog']:
            if base_path.lower().endswith(ext):
                base_path = base_path[:-len(ext)]
                break
        
        if self.use_frame_number:
            filepath = f"{base_path}_{frame_number:04d}{file_ext}"
        else:
            filepath = base_path + file_ext
        
        export_dir = os.path.dirname(filepath)

        try:
            axis_matrix = get_axis_conversion_matrix(self)

            all_samples = []
            for obj in selected_objects:
                if self.sampling_mode == 'VERTICES':
                    samples = sample_vertices(obj, context, axis_matrix, self)
                else:
                    samples = sample_mesh(obj, context, axis_matrix, self)
                all_samples.extend(samples)

            self.report({'INFO'}, f"Frame {frame_number}: Sampled {len(all_samples)} points from {len(selected_objects)} object(s)")

            # Choose export method
            if self.use_splat_transform:
                # Use splat-transform method
                base_name = os.path.splitext(os.path.basename(filepath))[0]
                generator_path = os.path.join(export_dir, f"{base_name}.mjs")
                
                create_mesh_generator(generator_path, all_samples)

                # Build splat-transform command
                cmd = [self.splat_transform_path]
                if self.overwrite_output:
                    cmd.append('-w')
                cmd.extend([generator_path, filepath])

                result = subprocess.run(
                    cmd,
                    capture_output=True,
                    text=True,
                    shell=True
                )

                if result.returncode != 0:
                    self.report({'ERROR'}, f"splat-transform failed: {result.stderr}")
                    return {'CANCELLED'}

                # Delete .mjs file if not keeping it
                if not self.keep_mjs_file and os.path.exists(generator_path):
                    try:
                        os.remove(generator_path)
                    except Exception as e:
                        self.report({'WARNING'}, f"Could not delete .mjs file: {str(e)}")
            else:
                # Write PLY directly
                if self.export_format == 'PLY':
                    write_ply_direct(filepath, all_samples, self.overwrite_output)
                else:
                    self.report({'ERROR'}, "Direct export only supports PLY format. Use splat-transform for SOG format.")
                    return {'CANCELLED'}

            self.report({'INFO'}, f"Successfully exported to {filepath}")
            return {'FINISHED'}

        except Exception as e:
            import traceback
            self.report({'ERROR'}, f"Export failed: {str(e)}")
            print(traceback.format_exc())  # Print full traceback to console
            return {'CANCELLED'}

    def batch_export_frames(self, context, selected_objects):
        """Export all frames in the timeline range"""
        start_frame = context.scene.frame_start
        end_frame = context.scene.frame_end
        original_frame = context.scene.frame_current
        
        total_frames = end_frame - start_frame + 1
        self.report({'INFO'}, f"Starting batch export of {total_frames} frames...")
        
        success_count = 0
        fail_count = 0
        
        for frame in range(start_frame, end_frame + 1):
            result = self.export_single_frame(context, selected_objects, frame)
            if result == {'FINISHED'}:
                success_count += 1
            else:
                fail_count += 1
        
        # Restore original frame
        context.scene.frame_set(original_frame)
        
        if fail_count == 0:
            self.report({'INFO'}, f"Batch export complete! Successfully exported {success_count} frames.")
            return {'FINISHED'}
        else:
            self.report({'WARNING'}, f"Batch export finished with errors. Success: {success_count}, Failed: {fail_count}")
            return {'FINISHED'}


# Shared utility functions
def get_axis_conversion_matrix(settings):
    """Create a conversion matrix based on axis settings"""
    from bpy_extras.io_utils import axis_conversion

    conv_matrix = axis_conversion(
        from_forward='-Y',
        from_up='Z',
        to_forward=settings.axis_forward,
        to_up=settings.axis_up,
    ).to_4x4()

    return conv_matrix


def sample_vertices(obj, context, axis_matrix, settings):
    """Create splats directly from mesh vertices"""
    depsgraph = context.evaluated_depsgraph_get()
    obj_eval = obj.evaluated_get(depsgraph)
    mesh = obj_eval.to_mesh()

    matrix_world = axis_matrix @ obj.matrix_world
    material = obj.active_material
    
    color_attribute = None
    has_vertex_colors = False
    
    if hasattr(mesh, 'color_attributes') and len(mesh.color_attributes) > 0:
        color_attribute = mesh.color_attributes.active_color
        if color_attribute:
            has_vertex_colors = True
    elif hasattr(mesh, 'vertex_colors') and len(mesh.vertex_colors) > 0:
        color_attribute = mesh.vertex_colors.active
        has_vertex_colors = True

    # Build KD-tree for nearest neighbor searches if auto-scale is enabled
    kd = None
    if settings.use_auto_scale:
        from mathutils import kdtree
        kd = kdtree.KDTree(len(mesh.vertices))
        for i, v in enumerate(mesh.vertices):
            kd.insert(v.co, i)
        kd.balance()

    samples = []

    for vert in mesh.vertices:
        world_pos = matrix_world @ vert.co
        color, alpha = get_vertex_color(vert.index, mesh, has_vertex_colors, color_attribute, material, settings)
        normal = matrix_world.to_3x3() @ vert.normal if settings.use_normals else Vector((0, 0, 1))
        
        # Calculate scale based on average nearest neighbor distance
        if settings.use_auto_scale and kd:
            # Find N+1 nearest vertices (first will be the vertex itself at distance 0)
            num_neighbors = settings.auto_scale_neighbors + 1
            nearest = kd.find_n(vert.co, num_neighbors)
            if len(nearest) > 1:
                # Skip the first result (self) and average the distances
                distances = [n[2] for n in nearest[1:]]  # n[2] is the distance
                avg_dist = sum(distances) / len(distances)
                # Use average distance as the scale (will give roughly correct coverage)
                auto_scale = max(avg_dist, 0.0001)
            else:
                auto_scale = 0.01  # Fallback for isolated vertices
            final_scale = auto_scale * settings.splat_scale
        else:
            final_scale = settings.splat_scale
        
        # Ensure final_scale is always positive and non-zero
        final_scale = max(final_scale, 0.0001)
        
        # Multiply alpha with global opacity
        final_opacity = alpha * settings.splat_opacity

        samples.append({
            'position': world_pos,
            'color': color,
            'normal': normal,
            'scale': final_scale,
            'opacity': final_opacity
        })

    obj_eval.to_mesh_clear()
    return samples


def sample_mesh(obj, context, axis_matrix, settings):
    """Sample points from mesh surface"""
    depsgraph = context.evaluated_depsgraph_get()
    obj_eval = obj.evaluated_get(depsgraph)
    mesh = obj_eval.to_mesh()

    matrix_world = axis_matrix @ obj.matrix_world
    bm = bmesh.new()
    bm.from_mesh(mesh)
    bm.faces.ensure_lookup_table()

    material = obj.active_material
    
    color_attribute = None
    has_vertex_colors = False
    
    if hasattr(mesh, 'color_attributes') and len(mesh.color_attributes) > 0:
        color_attribute = mesh.color_attributes.active_color
        if color_attribute:
            has_vertex_colors = True
    elif hasattr(mesh, 'vertex_colors') and len(mesh.vertex_colors) > 0:
        color_attribute = mesh.vertex_colors.active
        has_vertex_colors = True

    # Build KD-tree for nearest neighbor searches if auto-scale is enabled
    kd = None
    if settings.use_auto_scale:
        from mathutils import kdtree
        kd = kdtree.KDTree(len(mesh.vertices))
        for i, v in enumerate(mesh.vertices):
            kd.insert(v.co, i)
        kd.balance()

    samples = []

    total_area = sum(face.calc_area() for face in bm.faces)
    num_samples = int(total_area * settings.sample_density)

    for face in bm.faces:
        face_area = face.calc_area()
        face_samples = max(1, int((face_area / total_area) * num_samples))

        for _ in range(face_samples):
            r1 = math.sqrt(bpy.context.scene.frame_current * 0.001 + len(samples) * 0.1) % 1.0
            r2 = (len(samples) * 0.7) % 1.0
            if r1 + r2 > 1:
                r1 = 1 - r1
                r2 = 1 - r2

            v0, v1, v2 = [v.co for v in face.verts[:3]]
            pos = v0 + (v1 - v0) * r1 + (v2 - v0) * r2
            world_pos = matrix_world @ pos

            color, alpha = get_face_color(face, obj, mesh, has_vertex_colors, color_attribute, material, settings)
            normal = matrix_world.to_3x3() @ face.normal if settings.use_normals else Vector((0, 0, 1))
            
            # Calculate scale based on average nearest vertex distance
            if settings.use_auto_scale and kd:
                # Find N nearest vertices
                num_neighbors = settings.auto_scale_neighbors
                nearest_list = kd.find_n(pos, num_neighbors)
                if nearest_list:
                    # Average the distances to all found neighbors
                    distances = [n[2] for n in nearest_list]  # n[2] is the distance
                    avg_dist = sum(distances) / len(distances)
                    auto_scale = max(avg_dist, 0.0001)
                else:
                    auto_scale = 0.01  # Fallback
                final_scale = auto_scale * settings.splat_scale
            else:
                final_scale = settings.splat_scale
            
            # Ensure final_scale is always positive and non-zero
            final_scale = max(final_scale, 0.0001)
            
            # Multiply alpha with global opacity
            final_opacity = alpha * settings.splat_opacity

            samples.append({
                'position': world_pos,
                'color': color,
                'normal': normal,
                'scale': final_scale,
                'opacity': final_opacity
            })

    bm.free()
    obj_eval.to_mesh_clear()

    return samples

def get_face_color(face, obj, mesh, has_vertex_colors, color_attribute, material, settings):
    """Get color and alpha for a face"""
    if settings.use_vertex_colors and has_vertex_colors and color_attribute:
        try:
            colors = []
            alphas = []
            
            if hasattr(color_attribute, 'domain'):
                domain = color_attribute.domain
                
                if domain == 'CORNER':
                    for loop in face.loops:
                        mesh_loop_index = loop.index
                        if mesh_loop_index < len(color_attribute.data):
                            color_data = color_attribute.data[mesh_loop_index].color
                            colors.append(color_data[:3])
                            # Get alpha (4th component)
                            alphas.append(color_data[3] if len(color_data) > 3 else 1.0)
                
                elif domain == 'POINT':
                    for vert in face.verts:
                        vert_index = vert.index
                        if vert_index < len(color_attribute.data):
                            color_data = color_attribute.data[vert_index].color
                            colors.append(color_data[:3])
                            alphas.append(color_data[3] if len(color_data) > 3 else 1.0)
                
                elif domain == 'FACE':
                    face_index = face.index
                    if face_index < len(color_attribute.data):
                        color_data = color_attribute.data[face_index].color
                        alpha = color_data[3] if len(color_data) > 3 else 1.0
                        return list(color_data[:3]), alpha
            
            else:
                for loop_elem in mesh.loops:
                    for loop in face.loops:
                        if loop.index == loop_elem.index:
                            if loop_elem.index < len(color_attribute.data):
                                color_data = color_attribute.data[loop_elem.index].color
                                colors.append(color_data[:3])
                                alphas.append(color_data[3] if len(color_data) > 3 else 1.0)
                            break
            
            if colors:
                avg_color = [sum(c[i] for c in colors) / len(colors) for i in range(3)]
                avg_alpha = sum(alphas) / len(alphas) if alphas else 1.0
                return avg_color, avg_alpha
                
        except (IndexError, AttributeError):
            pass

    if material and material.use_nodes:
        for node in material.node_tree.nodes:
            if node.type == 'BSDF_PRINCIPLED':
                base_color = node.inputs['Base Color'].default_value
                alpha = base_color[3] if len(base_color) > 3 else 1.0
                return list(base_color[:3]), alpha

    return [1.0, 1.0, 1.0], 1.0


def get_vertex_color(vert_index, mesh, has_vertex_colors, color_attribute, material, settings):
    """Get color and alpha for a specific vertex"""
    if settings.use_vertex_colors and has_vertex_colors and color_attribute:
        try:
            if hasattr(color_attribute, 'domain'):
                domain = color_attribute.domain
                
                if domain == 'POINT':
                    if vert_index < len(color_attribute.data):
                        color_data = color_attribute.data[vert_index].color
                        alpha = color_data[3] if len(color_data) > 3 else 1.0
                        return list(color_data[:3]), alpha
                
                elif domain == 'CORNER':
                    colors = []
                    alphas = []
                    for loop in mesh.loops:
                        if loop.vertex_index == vert_index:
                            if loop.index < len(color_attribute.data):
                                color_data = color_attribute.data[loop.index].color
                                colors.append(color_data[:3])
                                alphas.append(color_data[3] if len(color_data) > 3 else 1.0)
                    if colors:
                        avg_color = [sum(c[i] for c in colors) / len(colors) for i in range(3)]
                        avg_alpha = sum(alphas) / len(alphas) if alphas else 1.0
                        return avg_color, avg_alpha
                
                elif domain == 'FACE':
                    colors = []
                    alphas = []
                    for poly in mesh.polygons:
                        if vert_index in poly.vertices:
                            if poly.index < len(color_attribute.data):
                                color_data = color_attribute.data[poly.index].color
                                colors.append(color_data[:3])
                                alphas.append(color_data[3] if len(color_data) > 3 else 1.0)
                    if colors:
                        avg_color = [sum(c[i] for c in colors) / len(colors) for i in range(3)]
                        avg_alpha = sum(alphas) / len(alphas) if alphas else 1.0
                        return avg_color, avg_alpha
                
        except (IndexError, AttributeError):
            pass

    if material and material.use_nodes:
        for node in material.node_tree.nodes:
            if node.type == 'BSDF_PRINCIPLED':
                base_color = node.inputs['Base Color'].default_value
                alpha = base_color[3] if len(base_color) > 3 else 1.0
                return list(base_color[:3]), alpha

    return [1.0, 1.0, 1.0], 1.0


def create_mesh_generator(path, samples):
    """Create a generator file for splat-transform"""
    def normal_to_quat(normal):
        up = Vector((0, 0, 1))
        if abs(normal.dot(up)) > 0.999:
            return [0, 0, 0, 1]
        axis = up.cross(normal).normalized()
        angle = math.acos(up.dot(normal))
        half_angle = angle / 2
        s = math.sin(half_angle)
        return [axis.x * s, axis.y * s, axis.z * s, math.cos(half_angle)]

    SH_C0 = 0.28209479177387814

    sample_data = []
    for sample in samples:
        pos = sample['position']
        color = sample['color']
        scale = math.log(sample['scale'])
        opacity = sample['opacity']
        quat = normal_to_quat(sample['normal'])
        sample_str = (
            f"[{pos.x:.6f}, {pos.y:.6f}, {pos.z:.6f}, "
            f"{scale:.6f}, {color[0]:.6f}, {color[1]:.6f}, {color[2]:.6f}, {opacity:.6f}, "
            f"{quat[0]:.6f}, {quat[1]:.6f}, {quat[2]:.6f}, {quat[3]:.6f}]"
        )
        sample_data.append(sample_str)

    samples_joined = ',\n            '.join(sample_data)

    code = f"""class Generator {{
    constructor() {{
        this.count = {len(samples)};
        
        this.columnNames = [
            'x', 'y', 'z',
            'scale_0', 'scale_1', 'scale_2',
            'f_dc_0', 'f_dc_1', 'f_dc_2', 'opacity',
            'rot_0', 'rot_1', 'rot_2', 'rot_3'
        ];
        
        const SH_C0 = {SH_C0};
        const packClr = (c) => (c - 0.5) / SH_C0;
        const packOpacity = (opacity) => (opacity <= 0) ? -20 : (opacity >= 1) ? 20 : -Math.log(1 / opacity - 1);
        
        const samples = [
            {samples_joined}
        ];
        
        this.getRow = (index, row) => {{
            const s = samples[index];
            row.x = s[0];
            row.y = s[1];
            row.z = s[2];
            row.scale_0 = s[3];
            row.scale_1 = s[3];
            row.scale_2 = s[3];
            row.f_dc_0 = packClr(s[4]);
            row.f_dc_1 = packClr(s[5]);
            row.f_dc_2 = packClr(s[6]);
            row.opacity = packOpacity(s[7]);
            row.rot_0 = s[8];
            row.rot_1 = s[9];
            row.rot_2 = s[10];
            row.rot_3 = s[11];
        }};
    }}
    static create(params) {{
        return new Generator();
    }}
}}
export {{ Generator }};
"""

    with open(path, 'w') as f:
        f.write(code)


def write_ply_direct(filepath, samples, overwrite):
    """Write Gaussian Splat PLY file directly without splat-transform"""
    import struct
    
    if not overwrite and os.path.exists(filepath):
        raise FileExistsError(f"File {filepath} already exists and overwrite is disabled")
    
    def normal_to_quat(normal):
        """Convert normal vector to quaternion rotation"""
        up = Vector((0, 0, 1))
        if abs(normal.dot(up)) > 0.999:
            return [0.0, 0.0, 0.0, 1.0]
        axis = up.cross(normal).normalized()
        angle = math.acos(max(-1.0, min(1.0, up.dot(normal))))
        half_angle = angle / 2
        s = math.sin(half_angle)
        return [axis.x * s, axis.y * s, axis.z * s, math.cos(half_angle)]
    
    # Spherical harmonics constant
    SH_C0 = 0.28209479177387814
    
    # Calculate properties for each splat
    splat_data = []
    for sample in samples:
        pos = sample['position']
        color = sample['color']
        normal = sample['normal']
        scale = sample['scale']
        opacity = sample['opacity']
        
        # Convert to Gaussian splat format
        # Position (x, y, z)
        x, y, z = pos.x, pos.y, pos.z
        
        # Normal to quaternion (for rotation)
        quat = normal_to_quat(normal)
        rot_0, rot_1, rot_2, rot_3 = quat
        
        # Scale (logarithmic, same for all 3 axes)
        log_scale = math.log(max(scale, 1e-8))
        scale_0 = scale_1 = scale_2 = log_scale
        
        # Color to spherical harmonics (DC component only)
        # Formula: (color - 0.5) / SH_C0
        f_dc_0 = (color[0] - 0.5) / SH_C0
        f_dc_1 = (color[1] - 0.5) / SH_C0
        f_dc_2 = (color[2] - 0.5) / SH_C0
        
        # Opacity (logit transform)
        if opacity <= 0:
            opacity_logit = -20.0
        elif opacity >= 1:
            opacity_logit = 20.0
        else:
            opacity_logit = -math.log(1.0 / opacity - 1.0)
        
        # Additional spherical harmonics (45 zeros for higher order SH)
        sh_rest = [0.0] * 45
        
        splat_data.append({
            'x': x, 'y': y, 'z': z,
            'nx': normal.x, 'ny': normal.y, 'nz': normal.z,
            'f_dc_0': f_dc_0, 'f_dc_1': f_dc_1, 'f_dc_2': f_dc_2,
            'sh_rest': sh_rest,
            'opacity': opacity_logit,
            'scale_0': scale_0, 'scale_1': scale_1, 'scale_2': scale_2,
            'rot_0': rot_0, 'rot_1': rot_1, 'rot_2': rot_2, 'rot_3': rot_3
        })
    
    # Write PLY file
    with open(filepath, 'wb') as f:
        # Write ASCII header
        header = f"""ply
format binary_little_endian 1.0
element vertex {len(splat_data)}
property float x
property float y
property float z
property float nx
property float ny
property float nz
property float f_dc_0
property float f_dc_1
property float f_dc_2
"""
        # Add f_rest_0 through f_rest_44 (45 additional SH coefficients)
        for i in range(45):
            header += f"property float f_rest_{i}\n"
        
        header += """property float opacity
property float scale_0
property float scale_1
property float scale_2
property float rot_0
property float rot_1
property float rot_2
property float rot_3
end_header
"""
        f.write(header.encode('ascii'))
        
        # Write binary vertex data
        for splat in splat_data:
            # Pack all properties as little-endian floats
            data = struct.pack('<fff', splat['x'], splat['y'], splat['z'])  # position
            data += struct.pack('<fff', splat['nx'], splat['ny'], splat['nz'])  # normal
            data += struct.pack('<fff', splat['f_dc_0'], splat['f_dc_1'], splat['f_dc_2'])  # DC SH
            # Add 45 higher-order SH coefficients (all zeros)
            for sh_val in splat['sh_rest']:
                data += struct.pack('<f', sh_val)
            data += struct.pack('<f', splat['opacity'])  # opacity
            data += struct.pack('<fff', splat['scale_0'], splat['scale_1'], splat['scale_2'])  # scale
            data += struct.pack('<ffff', splat['rot_0'], splat['rot_1'], splat['rot_2'], splat['rot_3'])  # rotation
            f.write(data)


def menu_func_export(self, context):
    self.layout.operator(GaussianSplatExporter.bl_idname, text="Gaussian Splat (.ply/.sog)")


# Registration
classes = (
    GaussianSplatSettings,
    GAUSSIANSPLAT_PT_MainPanel,
    GAUSSIANSPLAT_OT_DirectExport,
    GaussianSplatExporter,
)


def register():
    for cls in classes:
        bpy.utils.register_class(cls)
    
    bpy.types.Scene.gaussian_splat_settings = PointerProperty(type=GaussianSplatSettings)
    bpy.types.TOPBAR_MT_file_export.append(menu_func_export)


def unregister():
    bpy.types.TOPBAR_MT_file_export.remove(menu_func_export)
    
    del bpy.types.Scene.gaussian_splat_settings
    
    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)


if __name__ == "__main__":
    register()