test_editor.py

"""Tests for SplatForge editor operations."""

import sys
import os
import numpy as np
import pytest

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from splatforge.core.structures import SplatData
from splatforge.core.editor import SplatEditor
from splatforge.utils.helpers import create_sample_cloud


class TestFiltering:
    """Test filtering operations."""

    def setup_method(self):
        """Create test data."""
        self.data = create_sample_cloud(1000, 'sphere', seed=42)

    def test_filter_by_alpha(self):
        """Test alpha filtering."""
        result = SplatEditor.filter_by_alpha(self.data, 0.7, 1.0)
        assert result.num_splats < self.data.num_splats
        assert result.alphas.min() >= 0.7

    def test_filter_by_scale(self):
        """Test scale filtering."""
        result = SplatEditor.filter_by_scale(self.data, 0.01, 0.04)
        assert result.num_splats <= self.data.num_splats

    def test_filter_by_bounding_box(self):
        """Test bounding box filtering."""
        center = self.data.center
        result = SplatEditor.filter_by_bounding_box(
            self.data,
            center - 1.0,
            center + 1.0
        )
        assert result.num_splats <= self.data.num_splats

    def test_filter_by_distance(self):
        """Test distance filtering."""
        result = SplatEditor.filter_by_distance(self.data, self.data.center, 2.0)
        assert result.num_splats <= self.data.num_splats

    def test_random_subsample(self):
        """Test random subsampling."""
        result = SplatEditor.random_subsample(self.data, 100, seed=42)
        assert result.num_splats == 100

    def test_random_subsample_larger_than_input(self):
        """Test subsample with count larger than input."""
        result = SplatEditor.random_subsample(self.data, 10000)
        assert result.num_splats == self.data.num_splats


class TestSpatialOperations:
    """Test spatial operations."""

    def setup_method(self):
        self.data = create_sample_cloud(1000, 'cube', seed=42)

    def test_normalize_center(self):
        """Test center normalization."""
        result = SplatEditor.normalize_position(self.data, 'center')
        assert np.allclose(result.center, [0, 0, 0], atol=1e-5)

    def test_normalize_unit(self):
        """Test unit normalization."""
        result = SplatEditor.normalize_position(self.data, 'unit')
        extent = result.extent
        assert extent.max() <= 1.0 + 1e-5

    def test_transform_translation(self):
        """Test translation."""
        result = SplatEditor.transform(self.data, translation=[10, 20, 30])
        expected_center = self.data.center + np.array([10, 20, 30])
        assert np.allclose(result.center, expected_center, atol=1e-5)

    def test_transform_scale(self):
        """Test scaling."""
        result = SplatEditor.transform(self.data, scale=2.0)
        assert np.allclose(result.extent, self.data.extent * 2.0, atol=1e-4)

    def test_crop(self):
        """Test cropping."""
        center = self.data.center
        result = SplatEditor.crop_to_bounding_box(
            self.data,
            center - 0.5,
            center + 0.5
        )
        assert result.num_splats <= self.data.num_splats


class TestAttributeOperations:
    """Test attribute manipulation."""

    def setup_method(self):
        self.data = create_sample_cloud(500, 'sphere', seed=42)

    def test_adjust_brightness(self):
        """Test brightness adjustment."""
        result = SplatEditor.adjust_colors(self.data, brightness=2.0)
        assert result.colors.max() <= 1.0  # Should be clamped

    def test_adjust_scales(self):
        """Test scale adjustment."""
        result = SplatEditor.adjust_scales(self.data, multiplier=2.0)
        assert np.allclose(result.scales, self.data.scales * 2.0, atol=1e-6)

    def test_adjust_alphas(self):
        """Test alpha adjustment."""
        result = SplatEditor.adjust_alphas(self.data, multiplier=0.5)
        assert result.alphas.max() <= self.data.alphas.max()

    def test_set_uniform_color(self):
        """Test uniform color."""
        result = SplatEditor.set_uniform_color(self.data, (1.0, 0.0, 0.0))
        assert np.allclose(result.colors[:, 0], 1.0)
        assert np.allclose(result.colors[:, 1], 0.0)
        assert np.allclose(result.colors[:, 2], 0.0)

    def test_set_uniform_scale(self):
        """Test uniform scale."""
        result = SplatEditor.set_uniform_scale(self.data, 0.05)
        assert np.allclose(result.scales, 0.05)


class TestSplitAndMerge:
    """Test split and merge operations."""

    def test_split_by_grid(self):
        """Test grid splitting."""
        data = create_sample_cloud(1000, 'cube', seed=42)
        chunks = SplatEditor.split_by_grid(data, grid_size=2.0)
        assert len(chunks) > 1
        total = sum(c.num_splats for c in chunks)
        assert total == data.num_splats

    def test_merge(self):
        """Test merging."""
        d1 = create_sample_cloud(100, 'sphere', seed=42)
        d2 = create_sample_cloud(100, 'cube', seed=43)
        merged = SplatEditor.merge([d1, d2])
        assert merged.num_splats == 200

    def test_merge_empty_list(self):
        """Test merging empty list."""
        result = SplatEditor.merge([])
        assert result.is_empty


class TestQualityAnalysis:
    """Test quality analysis."""

    def test_clean_data_quality(self):
        """Test quality of clean data."""
        data = create_sample_cloud(1000, 'sphere', seed=42)
        quality = SplatEditor.analyze_quality(data)
        assert quality['score'] > 0
        assert isinstance(quality['issues'], list)
        assert isinstance(quality['warnings'], list)
        assert isinstance(quality['suggestions'], list)

    def test_auto_repair(self):
        """Test auto repair."""
        # Create data with issues
        positions = np.random.randn(100, 3).astype(np.float32)
        positions[0, 0] = float('nan')
        data = SplatData(
            positions=positions,
            scales=np.ones((100, 3), dtype=np.float32) * 0.01,
            rotations=np.tile([1, 0, 0, 0], (100, 1)).astype(np.float32),
            colors=np.random.rand(100, 3).astype(np.float32),
            alphas=np.random.rand(100, 1).astype(np.float32),
        )
        repaired, repairs = SplatEditor.auto_repair(data)
        assert repaired.num_splats < data.num_splats  # NaN point removed
        assert len(repairs) > 0


if __name__ == '__main__':
    pytest.main([__file__, '-v'])