"""Tests for compute_sdf_fast function - GPU-compatible SDF computation."""

import pytest
import torch
from sdf_loss.utils import compute_sdf, compute_sdf_fast


# Skip GPU tests if CUDA is not available
cuda_available = torch.cuda.is_available()
skip_no_cuda = pytest.mark.skipif(not cuda_available, reason="CUDA not available")


class TestComputeSDFFastBasic:
    """Basic functionality tests for compute_sdf_fast."""

    def test_basic_computation(self):
        """Test that compute_sdf_fast runs without errors on a simple mask."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32)
        mask[0, 0, 8:24, 8:24] = 1.0

        sdf = compute_sdf_fast(mask, normalize=True)

        assert sdf is not None, "Should return a result"
        assert torch.isfinite(sdf).all(), "All SDF values should be finite"

    def test_output_shape_4d(self):
        """Test output shape for 4D input (B, C, H, W)."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32)
        mask[0, 0, 8:24, 8:24] = 1.0

        sdf = compute_sdf_fast(mask, normalize=True)

        # Output should match input shape
        assert sdf.ndim == 4, f"Expected 4D output, got shape {sdf.shape}"
        assert sdf.shape == (1, 1, 32, 32), f"Expected (1, 1, 32, 32), got {sdf.shape}"

    def test_output_dtype(self):
        """Test that output dtype matches input dtype."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32)
        mask[0, 0, 8:24, 8:24] = 1.0

        sdf = compute_sdf_fast(mask, normalize=True)

        assert sdf.dtype == torch.float32, f"Expected float32, got {sdf.dtype}"

    def test_normalized_range(self):
        """Test that normalized SDF is in expected range."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32)
        mask[0, 0, 8:24, 8:24] = 1.0

        sdf = compute_sdf_fast(mask, normalize=True)

        assert sdf.min() >= -1.0, f"Min should be >= -1, got {sdf.min()}"
        assert sdf.max() <= 1.0, f"Max should be <= 1, got {sdf.max()}"

    def test_unnormalized(self):
        """Test unnormalized SDF computation."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32)
        mask[0, 0, 8:24, 8:24] = 1.0

        sdf_norm = compute_sdf_fast(mask, normalize=True)
        sdf_raw = compute_sdf_fast(mask, normalize=False)

        # Unnormalized should have larger magnitude
        assert sdf_raw.abs().max() >= sdf_norm.abs().max(), (
            "Raw SDF should have >= magnitude than normalized"
        )

    def test_positive_inside_negative_outside(self):
        """Test SDF sign convention: positive inside, negative outside."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32)
        mask[0, 0, 12:20, 12:20] = 1.0

        sdf = compute_sdf_fast(mask, normalize=False)

        # Center (inside) should be positive
        center_val = sdf[0, 0, 16, 16].item()
        assert center_val > 0, f"Center should be positive, got {center_val}"

        # Corner (outside) should be negative
        corner_val = sdf[0, 0, 0, 0].item()
        assert corner_val < 0, f"Corner should be negative, got {corner_val}"


class TestComputeSDFFastEdgeCases:
    """Edge case tests to identify failure modes."""

    def test_empty_mask(self):
        """Test SDF with all-zero mask (nothing is foreground)."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32)

        sdf = compute_sdf_fast(mask, normalize=True)

        # All pixels should be "outside" (negative or zero distance)
        assert (sdf <= 0).all(), "Empty mask should have all non-positive SDF"
        assert torch.isfinite(sdf).all(), "Should not produce NaN/Inf on empty mask"

    def test_full_mask(self):
        """Test SDF with all-one mask (everything is foreground)."""
        mask = torch.ones(1, 1, 32, 32, dtype=torch.float32)

        sdf = compute_sdf_fast(mask, normalize=True)

        # All pixels should be "inside" (positive or zero distance)
        assert (sdf >= 0).all(), "Full mask should have all non-negative SDF"
        assert torch.isfinite(sdf).all(), "Should not produce NaN/Inf on full mask"

    def test_single_pixel_foreground(self):
        """Test SDF with single foreground pixel."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32)
        mask[0, 0, 16, 16] = 1.0

        sdf = compute_sdf_fast(mask, normalize=True)

        assert torch.isfinite(sdf).all(), "Should handle single pixel foreground"
        # Single pixel is "inside" the foreground region, so positive
        # After normalization it becomes 1.0 (max positive value)
        assert sdf[0, 0, 16, 16] > 0, "Single foreground pixel should have positive SDF"
        # Surrounding pixels should be negative (outside)
        assert sdf[0, 0, 0, 0] < 0, "Background pixels should have negative SDF"

    def test_single_pixel_background(self):
        """Test SDF with single background pixel in full mask."""
        mask = torch.ones(1, 1, 32, 32, dtype=torch.float32)
        mask[0, 0, 16, 16] = 0.0

        sdf = compute_sdf_fast(mask, normalize=True)

        assert torch.isfinite(sdf).all(), "Should handle single pixel background"

    def test_thin_line_horizontal(self):
        """Test SDF with thin horizontal line (1 pixel wide)."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32)
        mask[0, 0, 16, :] = 1.0

        sdf = compute_sdf_fast(mask, normalize=True)

        assert torch.isfinite(sdf).all(), "Should handle thin horizontal line"

    def test_thin_line_vertical(self):
        """Test SDF with thin vertical line (1 pixel wide)."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32)
        mask[0, 0, :, 16] = 1.0

        sdf = compute_sdf_fast(mask, normalize=True)

        assert torch.isfinite(sdf).all(), "Should handle thin vertical line"

    def test_checkerboard_pattern(self):
        """Test SDF with alternating checkerboard pattern."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32)
        mask[0, 0, ::2, ::2] = 1.0
        mask[0, 0, 1::2, 1::2] = 1.0

        sdf = compute_sdf_fast(mask, normalize=True)

        assert torch.isfinite(sdf).all(), "Should handle checkerboard pattern"

    def test_small_image(self):
        """Test SDF with very small image."""
        mask = torch.zeros(1, 1, 4, 4, dtype=torch.float32)
        mask[0, 0, 1:3, 1:3] = 1.0

        sdf = compute_sdf_fast(mask, normalize=True)

        assert sdf is not None, "Should handle small images"
        assert torch.isfinite(sdf).all(), (
            "Should produce finite values for small images"
        )

    def test_non_square_image(self):
        """Test SDF with non-square image."""
        mask = torch.zeros(1, 1, 64, 32, dtype=torch.float32)
        mask[0, 0, 20:44, 8:24] = 1.0

        sdf = compute_sdf_fast(mask, normalize=True)

        assert sdf.shape == (1, 1, 64, 32), f"Should handle non-square, got {sdf.shape}"
        assert torch.isfinite(sdf).all(), "Should produce finite values for non-square"


class TestComputeSDFFastInputValidation:
    """Input validation and error handling tests."""

    def test_rejects_values_above_one(self):
        """Test that values > 1.0 are rejected."""
        mask = torch.ones(1, 1, 32, 32, dtype=torch.float32) * 1.5

        with pytest.raises(AssertionError):
            compute_sdf_fast(mask, normalize=True)

    def test_rejects_negative_values(self):
        """Test that negative values are rejected."""
        mask = torch.ones(1, 1, 32, 32, dtype=torch.float32) * -0.5

        with pytest.raises(AssertionError):
            compute_sdf_fast(mask, normalize=True)

    def test_accepts_soft_masks(self):
        """Test that soft/probabilistic masks (values between 0-1) work."""
        mask = torch.rand(1, 1, 32, 32, dtype=torch.float32)

        sdf = compute_sdf_fast(mask, normalize=True)

        assert torch.isfinite(sdf).all(), "Should handle soft masks"

    def test_2d_input(self):
        """Test behavior with 2D input (H, W) - may need reshaping."""
        mask = torch.zeros(32, 32, dtype=torch.float32)
        mask[8:24, 8:24] = 1.0

        # Document expected behavior: should it work or require 4D?
        try:
            sdf = compute_sdf_fast(mask.unsqueeze(0).unsqueeze(0), normalize=True)
            assert torch.isfinite(sdf).all()
        except Exception as e:
            pytest.fail(f"4D reshaping should work: {e}")

    def test_3d_input(self):
        """Test behavior with 3D input (B, H, W)."""
        mask = torch.zeros(2, 32, 32, dtype=torch.float32)
        mask[0, 8:24, 8:24] = 1.0
        mask[1, 4:28, 4:28] = 1.0

        # Document expected behavior
        try:
            # Try with unsqueeze to make 4D
            result = compute_sdf_fast(mask.unsqueeze(1), normalize=True)
            assert torch.isfinite(result).all()
        except Exception:
            # This is expected if function only handles single image
            pass


class TestComputeSDFFastDeviceHandling:
    """Tests for device (CPU/GPU) handling."""

    def test_cpu_input_stays_on_cpu(self):
        """Test that CPU input produces CPU output."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32, device="cpu")
        mask[0, 0, 8:24, 8:24] = 1.0

        sdf = compute_sdf_fast(mask, normalize=True)

        assert sdf.device.type == "cpu", f"Expected CPU output, got {sdf.device}"

    @skip_no_cuda
    def test_gpu_input_stays_on_gpu(self):
        """Test that GPU input produces GPU output without unnecessary transfers."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32, device="cuda")
        mask[0, 0, 8:24, 8:24] = 1.0

        sdf = compute_sdf_fast(mask, normalize=True)

        assert sdf.device.type == "cuda", f"Expected CUDA output, got {sdf.device}"

    @skip_no_cuda
    def test_no_unnecessary_cpu_transfers(self):
        """Test that GPU computation doesn't silently transfer to CPU."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32, device="cuda")
        mask[0, 0, 8:24, 8:24] = 1.0

        # Record CUDA memory before
        torch.cuda.synchronize()

        sdf = compute_sdf_fast(mask, normalize=True)

        torch.cuda.synchronize()

        # Output should be on GPU
        assert sdf.is_cuda, "Output should remain on GPU"
        # Input should not have been modified
        assert mask.is_cuda, "Input should remain on GPU"

    @skip_no_cuda
    def test_gpu_computation_actually_runs_on_gpu(self):
        """Verify computation happens on GPU, not secretly on CPU."""
        mask = torch.zeros(1, 1, 64, 64, dtype=torch.float32, device="cuda")
        mask[0, 0, 16:48, 16:48] = 1.0

        # Clear cache
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

        sdf = compute_sdf_fast(mask, normalize=True)

        torch.cuda.synchronize()

        # Result must be on GPU
        assert sdf.is_cuda, "Result must be on GPU"


class TestComputeSDFFastTraining:
    """Tests for training scenarios - gradients and autograd compatibility."""

    def test_gradient_flow(self):
        """Test that gradients flow through compute_sdf_fast."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32, requires_grad=True)
        # Can't use in-place ops with requires_grad, so create differently
        mask_data = torch.zeros(1, 1, 32, 32, dtype=torch.float32)
        mask_data[0, 0, 8:24, 8:24] = 1.0
        mask = mask_data.clone().requires_grad_(True)

        sdf = compute_sdf_fast(mask, normalize=True)

        # Try to backprop
        if sdf.requires_grad:
            loss = sdf.sum()
            loss.backward()
            assert mask.grad is not None, "Gradients should flow through"
            assert torch.isfinite(mask.grad).all(), "Gradients should be finite"
        else:
            # Document that gradients don't flow (FastGeodis may not be differentiable)
            pytest.skip("compute_sdf_fast does not support gradient flow")

    def test_no_inplace_modification(self):
        """Test that input tensor is not modified in-place."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32)
        mask[0, 0, 8:24, 8:24] = 1.0
        mask_original = mask.clone()

        _ = compute_sdf_fast(mask, normalize=True)

        assert torch.allclose(mask, mask_original), "Input should not be modified"

    @skip_no_cuda
    def test_gradient_flow_gpu(self):
        """Test gradient flow on GPU."""
        mask_data = torch.zeros(1, 1, 32, 32, dtype=torch.float32, device="cuda")
        mask_data[0, 0, 8:24, 8:24] = 1.0
        mask = mask_data.clone().requires_grad_(True)

        sdf = compute_sdf_fast(mask, normalize=True)

        if sdf.requires_grad:
            loss = sdf.sum()
            loss.backward()
            assert mask.grad is not None, "GPU gradients should flow through"
            assert mask.grad.is_cuda, "Gradients should be on GPU"
        else:
            pytest.skip("compute_sdf_fast does not support gradient flow on GPU")

    def test_batch_processing(self):
        """Test processing multiple samples in a batch."""
        batch_size = 4
        masks = torch.zeros(batch_size, 1, 32, 32, dtype=torch.float32)
        for i in range(batch_size):
            offset = i * 2
            masks[i, 0, 8 + offset : 24 - offset, 8 + offset : 24 - offset] = 1.0

        sdf = compute_sdf_fast(masks, normalize=True)

        # Verify output shape matches input
        assert sdf.shape == masks.shape, f"Expected {masks.shape}, got {sdf.shape}"
        assert torch.isfinite(sdf).all(), "All values should be finite"

        # Verify each sample has correct sign convention
        for i in range(batch_size):
            offset = i * 2
            center = 16
            # Center should be positive (inside)
            assert sdf[i, 0, center, center] > 0, (
                f"Sample {i} center should be positive"
            )
            # Corner should be negative (outside)
            assert sdf[i, 0, 0, 0] < 0, f"Sample {i} corner should be negative"

    def test_multi_channel(self):
        """Test processing multiple channels."""
        masks = torch.zeros(2, 3, 32, 32, dtype=torch.float32)
        # Different patterns per channel
        masks[0, 0, 8:24, 8:24] = 1.0
        masks[0, 1, 4:28, 4:28] = 1.0
        masks[0, 2, 12:20, 12:20] = 1.0
        masks[1, 0, 10:22, 10:22] = 1.0
        masks[1, 1, 6:26, 6:26] = 1.0
        masks[1, 2, 14:18, 14:18] = 1.0

        sdf = compute_sdf_fast(masks, normalize=True)

        # Verify output shape matches input
        assert sdf.shape == (2, 3, 32, 32), f"Expected (2, 3, 32, 32), got {sdf.shape}"
        assert torch.isfinite(sdf).all(), "All values should be finite"

        # Each channel should be processed independently
        for b in range(2):
            for c in range(3):
                # Center should be positive
                assert sdf[b, c, 16, 16] > 0, (
                    f"Batch {b} channel {c} center should be positive"
                )
                # Corner should be negative
                assert sdf[b, c, 0, 0] < 0, (
                    f"Batch {b} channel {c} corner should be negative"
                )


class TestComputeSDFFastConsistency:
    """Tests comparing compute_sdf_fast with compute_sdf for consistency."""

    def test_similar_results_to_compute_sdf(self):
        """Test that compute_sdf_fast produces similar results to compute_sdf."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32)
        mask[0, 0, 8:24, 8:24] = 1.0

        sdf_fast = compute_sdf_fast(mask, normalize=True)
        sdf_orig = compute_sdf(mask, normalize=True)

        # Results may not be identical due to different algorithms
        # but should have same sign pattern and similar magnitudes

        # Check same sign at key points
        assert (sdf_fast[0, 0, 16, 16] > 0) == (sdf_orig[0, 0, 16, 16] > 0), (
            "Center sign should match"
        )
        assert (sdf_fast[0, 0, 0, 0] < 0) == (sdf_orig[0, 0, 0, 0] < 0), (
            "Corner sign should match"
        )

    def test_deterministic(self):
        """Test that compute_sdf_fast is deterministic."""
        mask = torch.zeros(1, 1, 32, 32, dtype=torch.float32)
        mask[0, 0, 8:24, 8:24] = 1.0

        sdf1 = compute_sdf_fast(mask, normalize=True)
        sdf2 = compute_sdf_fast(mask, normalize=True)

        assert torch.allclose(sdf1, sdf2), "Results should be deterministic"


class TestComputeSDFFastPerformance:
    """Performance-related tests."""

    def test_large_image(self):
        """Test that large images work without OOM."""
        mask = torch.zeros(1, 1, 512, 512, dtype=torch.float32)
        mask[0, 0, 128:384, 128:384] = 1.0

        sdf = compute_sdf_fast(mask, normalize=True)

        assert sdf.shape == (1, 1, 512, 512), "Should handle large images"
        assert torch.isfinite(sdf).all(), "Large image should produce finite results"

    @skip_no_cuda
    def test_large_image_gpu(self):
        """Test large images on GPU."""
        mask = torch.zeros(1, 1, 512, 512, dtype=torch.float32, device="cuda")
        mask[0, 0, 128:384, 128:384] = 1.0

        sdf = compute_sdf_fast(mask, normalize=True)

        assert sdf.is_cuda, "Should remain on GPU"
        assert torch.isfinite(sdf).all(), "Should produce finite results on GPU"


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