"""
test_depth_hole_fill.py — Unit tests for depth hole fill helpers (no ROS2 required).

Covers:
  valid_mask:
    - valid range returns True
    - zero / below d_min returns False
    - NaN returns False
    - above d_max returns False
    - mixed array has correct mask

  _odd_kernel_size:
    - result is always odd
    - result >= 3
    - scales correctly

  fill_holes — no-hole cases:
    - fully valid image is returned unchanged
    - output dtype is float32
    - output shape matches input

  fill_holes — basic fills:
    - single centre hole in uniform depth → filled with correct depth
    - single centre hole in uniform depth → original valid pixels unchanged
    - NaN pixel treated as hole and filled
    - row of zeros within uniform depth → filled

  fill_holes — fill quality:
    - linear gradient: centre hole filled with interpolated value
    - multi-pass fills larger holes than single pass
    - all-zero image stays zero (no valid neighbours)
    - border hole (edge pixel) is handled without crash
    - depth range: pixel above d_max treated as hole

  fill_holes — valid pixel preservation:
    - original valid pixels are never modified
    - max_passes=1 still fills small holes
"""

import sys
import os
import math

import numpy as np
import pytest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from saltybot_bringup._depth_hole_fill import (
    fill_holes,
    valid_mask,
    _odd_kernel_size,
)


# ── Helpers ───────────────────────────────────────────────────────────────────

def _uniform(val=2.0, h=64, w=64) -> np.ndarray:
    return np.full((h, w), val, dtype=np.float32)


def _poke_hole(arr, r, c) -> np.ndarray:
    arr = arr.copy()
    arr[r, c] = 0.0
    return arr


def _poke_nan(arr, r, c) -> np.ndarray:
    arr = arr.copy()
    arr[r, c] = float('nan')
    return arr


# ── valid_mask ────────────────────────────────────────────────────────────────

class TestValidMask:

    def test_valid_pixel_is_true(self):
        d = np.array([[1.0]], dtype=np.float32)
        assert valid_mask(d, 0.1, 10.0)[0, 0]

    def test_zero_is_false(self):
        d = np.array([[0.0]], dtype=np.float32)
        assert not valid_mask(d, 0.1, 10.0)[0, 0]

    def test_below_dmin_is_false(self):
        d = np.array([[0.05]], dtype=np.float32)
        assert not valid_mask(d, 0.1, 10.0)[0, 0]

    def test_nan_is_false(self):
        d = np.array([[float('nan')]], dtype=np.float32)
        assert not valid_mask(d, 0.1, 10.0)[0, 0]

    def test_above_dmax_is_false(self):
        d = np.array([[15.0]], dtype=np.float32)
        assert not valid_mask(d, 0.1, 10.0)[0, 0]

    def test_at_dmin_is_true(self):
        d = np.array([[0.1]], dtype=np.float32)
        assert valid_mask(d, 0.1, 10.0)[0, 0]

    def test_at_dmax_is_true(self):
        d = np.array([[10.0]], dtype=np.float32)
        assert valid_mask(d, 0.1, 10.0)[0, 0]

    def test_mixed_array(self):
        d = np.array([[0.0, 1.0, float('nan'), 5.0, 11.0]], dtype=np.float32)
        m = valid_mask(d, 0.1, 10.0)
        np.testing.assert_array_equal(m, [[False, True, False, True, False]])


# ── _odd_kernel_size ──────────────────────────────────────────────────────────

class TestOddKernelSize:

    @pytest.mark.parametrize('base,scale', [
        (5, 1.0), (5, 2.5), (5, 6.0),
        (3, 1.0), (7, 2.0), (9, 3.0),
        (4, 1.0),   # even base → must become odd
    ])
    def test_result_is_odd(self, base, scale):
        ks = _odd_kernel_size(base, scale)
        assert ks % 2 == 1

    @pytest.mark.parametrize('base,scale', [(3, 1.0), (1, 5.0), (2, 0.5)])
    def test_result_at_least_3(self, base, scale):
        assert _odd_kernel_size(base, scale) >= 3

    def test_scale_1_returns_base_or_nearby_odd(self):
        ks = _odd_kernel_size(5, 1.0)
        assert ks == 5

    def test_large_scale_gives_large_kernel(self):
        ks = _odd_kernel_size(5, 6.0)
        assert ks >= 25   # 5 * 6 = 30 → 31


# ── fill_holes — output contract ──────────────────────────────────────────────

class TestFillHolesOutputContract:

    def test_output_dtype_float32(self):
        out = fill_holes(_uniform(2.0))
        assert out.dtype == np.float32

    def test_output_shape_preserved(self):
        img = _uniform(2.0, h=48, w=64)
        out = fill_holes(img)
        assert out.shape == img.shape

    def test_fully_valid_image_unchanged(self):
        img = _uniform(2.0)
        out = fill_holes(img)
        np.testing.assert_allclose(out, img, atol=1e-6)

    def test_valid_pixels_never_modified(self):
        """Any pixel valid in the input must be identical in the output."""
        img = _uniform(3.0, h=32, w=32)
        img[16, 16] = 0.0   # one hole
        mask_before = valid_mask(img)
        out = fill_holes(img)
        np.testing.assert_allclose(out[mask_before], img[mask_before], atol=1e-6)


# ── fill_holes — basic hole filling ──────────────────────────────────────────

class TestFillHolesBasic:

    def test_centre_zero_filled_uniform(self):
        """Single zero pixel in uniform depth → filled with that depth."""
        img = _poke_hole(_uniform(2.0, 32, 32), 16, 16)
        out = fill_holes(img, kernel_size=5, max_passes=1)
        assert out[16, 16] == pytest.approx(2.0, abs=0.05)

    def test_centre_nan_filled_uniform(self):
        """Single NaN pixel in uniform depth → filled."""
        img = _poke_nan(_uniform(2.0, 32, 32), 16, 16)
        out = fill_holes(img, kernel_size=5, max_passes=1)
        assert out[16, 16] == pytest.approx(2.0, abs=0.05)

    def test_filled_value_is_positive(self):
        img = _poke_hole(_uniform(1.5, 32, 32), 16, 16)
        out = fill_holes(img)
        assert out[16, 16] > 0.0

    def test_row_of_holes_filled(self):
        """Entire middle row zeroed → should be filled from neighbours above/below."""
        img = _uniform(3.0, 32, 32)
        img[16, :] = 0.0
        out = fill_holes(img, kernel_size=7, max_passes=1)
        # All pixels in the row should be non-zero after filling
        assert (out[16, :] > 0.0).all()

    def test_all_zero_stays_zero(self):
        """Image with no valid pixels → stays zero (nothing to interpolate from)."""
        img = np.zeros((32, 32), dtype=np.float32)
        out = fill_holes(img, d_min=0.1)
        assert (out == 0.0).all()

    def test_border_hole_no_crash(self):
        """Holes at image corners must not raise exceptions."""
        img = _uniform(2.0, 32, 32)
        img[0, 0] = 0.0
        img[0, -1] = 0.0
        img[-1, 0] = 0.0
        img[-1, -1] = 0.0
        out = fill_holes(img)   # must not raise
        assert out.shape == img.shape

    def test_border_holes_filled(self):
        """Corner holes should be filled from their neighbours."""
        img = _uniform(2.0, 32, 32)
        img[0, 0] = 0.0
        out = fill_holes(img, kernel_size=5, max_passes=1)
        assert out[0, 0] == pytest.approx(2.0, abs=0.1)


# ── fill_holes — fill quality ─────────────────────────────────────────────────

class TestFillHolesQuality:

    def test_linear_gradient_centre_hole_interpolated(self):
        """
        Depth linearly increasing from 1.0 (left) to 3.0 (right).
        Centre hole should be filled near the midpoint (~2.0).
        """
        h, w = 32, 32
        img = np.tile(np.linspace(1.0, 3.0, w, dtype=np.float32), (h, 1))
        cx = w // 2
        img[:, cx] = 0.0
        out = fill_holes(img, kernel_size=5, max_passes=1)
        mid = out[h // 2, cx]
        assert 1.5 <= mid <= 2.5, f'interpolated value {mid:.3f} not in [1.5, 2.5]'

    def test_large_hole_filled_with_more_passes(self):
        """A 9×9 hole in uniform depth: single pass may not fully fill it,
        but 3 passes should."""
        img = _uniform(2.0, 64, 64)
        # Create a 9×9 hole
        img[28:37, 28:37] = 0.0
        out1 = fill_holes(img, kernel_size=5, max_passes=1)
        out3 = fill_holes(img, kernel_size=5, max_passes=3)
        # More passes → fewer remaining holes
        holes1 = (out1 == 0.0).sum()
        holes3 = (out3 == 0.0).sum()
        assert holes3 <= holes1, f'more passes should reduce holes: {holes3} vs {holes1}'

    def test_3pass_fills_9x9_hole_completely(self):
        img = _uniform(2.0, 64, 64)
        img[28:37, 28:37] = 0.0
        out = fill_holes(img, kernel_size=5, max_passes=3)
        assert (out[28:37, 28:37] > 0.0).all()

    def test_filled_depth_within_valid_range(self):
        """Filled pixels should have depth within [d_min, d_max]."""
        img = _uniform(2.0, 32, 32)
        img[10:15, 10:15] = 0.0
        out = fill_holes(img, d_min=0.1, d_max=10.0, max_passes=3)
        # Only check pixels that were actually filled
        was_hole = (img == 0.0)
        filled   = out[was_hole]
        positive = filled[filled > 0.0]
        assert (positive >= 0.1).all()
        assert (positive <= 10.0).all()

    def test_above_dmax_treated_as_hole(self):
        """Pixels above d_max should be treated as holes and filled."""
        img = _uniform(2.0, 32, 32)
        img[16, 16] = 15.0   # out of range
        out = fill_holes(img, d_max=10.0, max_passes=1)
        assert out[16, 16] == pytest.approx(2.0, abs=0.1)

    def test_max_passes_1_works(self):
        img = _poke_hole(_uniform(2.0, 32, 32), 16, 16)
        out = fill_holes(img, max_passes=1)
        assert out.shape == img.shape
        assert out[16, 16] > 0.0


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