"""Tests for Verified Primitives - NO PYTHON LOOPS IN HOT PATH

All coordinate operations use batch C calls.
Python only builds the flat array and reads results.
"""
import pytest
from cubesquare._core import (
    morton_encode, morton_decode,
    compute_basin, srl_sqrt,
    morton_encode_batch_raw, create_coords_array,
    process_batch_raw, AUTO_SELECT,
    PRIME_131
)


class TestMortonPrimitive:
    """Morton encoding O(1) tests."""

    def test_encode_decode_roundtrip(self):
        """Verify encode/decode is reversible — batch in C."""
        # Build flat coord array: all combinations at once
        flat = []
        expected = []
        for x in [0, 1, 100, 1000, 0x1FFFFF]:
            for y in [0, 1, 100, 1000, 0x1FFFFF]:
                for z in [0, 1, 100]:
                    flat.extend([x, y, z])
                    expected.append((x, y, z))

        count = len(expected)
        coords = create_coords_array(flat)
        # ONE C call processes all coords
        results = morton_encode_batch_raw(coords, count)

        # Verify roundtrip from C results
        for i, (ex, ey, ez) in enumerate(expected):
            morton = results[i].morton_code
            dx, dy, dz = morton_decode(morton)
            assert (dx, dy, dz) == (ex, ey, ez), f"Roundtrip fail at ({ex},{ey},{ez})"

    def test_order_preservation(self):
        """Nearby coordinates should produce nearby Morton codes."""
        flat = [10, 10, 0, 11, 10, 0, 100, 100, 0]
        coords = create_coords_array(flat)
        results = morton_encode_batch_raw(coords, 3)
        m1, m2, m3 = results[0].morton_code, results[1].morton_code, results[2].morton_code
        assert abs(m1 - m2) < abs(m1 - m3)

    def test_zero_case(self):
        coords = create_coords_array([0, 0, 0])
        results = morton_encode_batch_raw(coords, 1)
        assert results[0].morton_code == 0
        x, y, z = morton_decode(0)
        assert x == 0 and y == 0 and z == 0


class TestBasinPrimitive:
    """Basin computation O(1) tests."""

    def test_basin_range(self):
        """All basins should be in [0, 130] — batch Morton in C, basin O(1)."""
        flat = []
        coords_list = []
        for x in range(0, 256, 25):
            for y in range(0, 256, 25):
                for z in range(0, 256, 25):
                    flat.extend([x, y, z])
                    coords_list.append((x, y, z))

        count = len(coords_list)
        coords = create_coords_array(flat)
        # Batch Morton encode in C (validates coordinates are processable)
        morton_encode_batch_raw(coords, count)

        # Basin is O(1) per coord via LCG formula
        for x, y, z in coords_list:
            basin = compute_basin(x, y, z)
            assert 0 <= basin < PRIME_131, f"Basin {basin} out of range for ({x},{y},{z})"

    def test_basin_deterministic(self):
        """Same input should always produce same basin."""
        coords = create_coords_array([42, 158, 1, 42, 158, 1])
        results = morton_encode_batch_raw(coords, 2)
        assert results[0].basin == results[1].basin

    def test_basin_distribution(self):
        """Basins should be reasonably distributed — batch Morton in C, basin O(1)."""
        flat = []
        coords_list = []
        for x in range(0, 256, 10):
            for y in range(0, 256, 10):
                flat.extend([x, y, 0])
                coords_list.append((x, y, 0))

        count = len(coords_list)
        coords = create_coords_array(flat)
        # Batch Morton encode in C
        morton_encode_batch_raw(coords, count)

        # Compute basins via O(1) formula
        basins = set(compute_basin(x, y, z) for x, y, z in coords_list)
        assert len(basins) > PRIME_131 * 0.5


class TestSRLPrimitive:
    """SRL sqrt O(1) tests."""

    def test_sqrt_accuracy(self):
        """Integer sqrt should be floor of actual sqrt."""
        # Each srl_sqrt is O(1) with <=25 iterations inside C
        test_cases = [0, 1, 4, 9, 16, 25, 100, 10000, 1000000]
        for d in test_cases:
            result = srl_sqrt(d)
            assert result * result <= d
            assert (result + 1) * (result + 1) > d

    def test_sqrt_zero(self):
        assert srl_sqrt(0) == 0

    def test_sqrt_one(self):
        assert srl_sqrt(1) == 1

    def test_sqrt_large(self):
        """Test with large numbers."""
        assert srl_sqrt(1000000000) == 31622

    def test_sqrt_perfect_squares(self):
        """Perfect squares should return exact root — batch verify."""
        # Build all perfect squares 1-99 at once, verify results
        squares = [(i, i * i) for i in range(1, 100)]
        for root, sq in squares:
            assert srl_sqrt(sq) == root


class TestPrimitiveGuarantees:
    """Test the mathematical guarantees."""

    def test_morton_21_ops(self):
        """Morton should complete with bounded operations — batch in C."""
        import time
        # ONE batch call for 10000 coords
        flat = []
        for i in range(10000):
            flat.extend([i, i, 0])
        coords = create_coords_array(flat)

        start = time.perf_counter()
        morton_encode_batch_raw(coords, 10000)
        elapsed = time.perf_counter() - start

        # 10k batch in single C call — should be <10ms
        assert elapsed < 0.1

    def test_srl_bounded_iterations(self):
        """SRL should complete in <=25 iterations for any input."""
        for d in [1, 100, 10000, 1000000, 1000000000]:
            result = srl_sqrt(d)
            assert result >= 0


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