"""Unit tests for the Gabor filter bank and Difference of Gaussians."""

import numpy as np
import pytest

from gabor import (
    gabor_kernel_2d,
    gabor_bank,
    gabor_response,
    gabor_feature_stack,
    dominant_orientation,
    dog,
)


def _grating(n=128, frequency=0.1, theta=0.0):
    """A sinusoidal grating of given spatial frequency and orientation."""
    y, x = np.mgrid[0:n, 0:n]
    xr = x * np.cos(theta) + y * np.sin(theta)
    return np.cos(2 * np.pi * frequency * xr)


# ---------------------------------------------------------------------------
# Single kernel
# ---------------------------------------------------------------------------

class TestKernel:
    def test_is_complex(self):
        k = gabor_kernel_2d(0.1, theta=0.0)
        assert np.iscomplexobj(k)

    def test_real_part_is_dc_free(self):
        # A bandpass filter has (near) zero mean: it rejects DC.
        k = gabor_kernel_2d(0.1, theta=0.0, bandwidth=1.0)
        assert abs(k.real.sum()) < 1e-2 * np.abs(k.real).sum()

    def test_orientation_changes_kernel(self):
        k0 = gabor_kernel_2d(0.1, theta=0.0)
        k90 = gabor_kernel_2d(0.1, theta=np.pi / 2)
        # Different orientations produce genuinely different kernels.
        assert k0.shape == k90.shape
        assert not np.allclose(k0, k90)


# ---------------------------------------------------------------------------
# Orientation selectivity
# ---------------------------------------------------------------------------

class TestSelectivity:
    def test_grating_excites_matching_orientation(self):
        f = 0.1
        # Vertical bars (grating varies along x) => theta = 0 responds most.
        img = _grating(frequency=f, theta=0.0)
        orientations = np.deg2rad([0, 45, 90, 135])
        best = dominant_orientation(img, f, orientations)
        assert best == pytest.approx(0.0)

    def test_rotated_grating_tracks_orientation(self):
        f = 0.1
        img = _grating(frequency=f, theta=np.pi / 2)  # horizontal bars
        orientations = np.deg2rad([0, 45, 90, 135])
        best = dominant_orientation(img, f, orientations)
        assert best == pytest.approx(np.pi / 2)

    def test_offaxis_orientation_within_one_grid_step(self):
        # A grating at 37deg lies strictly between the 15deg grid points, so
        # this checks near-match resolution, not just exact-grid recovery.
        f = 0.1
        true = np.deg2rad(37)
        img = _grating(frequency=f, theta=true)
        orientations = np.deg2rad(np.arange(0, 180, 15))
        best = dominant_orientation(img, f, orientations)
        assert abs(np.rad2deg(best) - 37) <= 15

    def test_response_is_nonnegative(self):
        img = _grating()
        r = gabor_response(img, 0.1, theta=0.0)
        assert np.all(r >= 0)


# ---------------------------------------------------------------------------
# Bank and feature stack
# ---------------------------------------------------------------------------

class TestBank:
    def test_bank_size(self):
        freqs = [0.1, 0.2, 0.3]
        thetas = np.deg2rad([0, 45, 90, 135])
        bank = gabor_bank(freqs, thetas)
        assert len(bank) == len(freqs) * len(thetas)

    def test_bank_entries_have_kernels(self):
        bank = gabor_bank([0.15], np.deg2rad([0, 90]))
        assert all(np.iscomplexobj(b["kernel"]) for b in bank)

    def test_feature_stack_shape(self):
        img = _grating(n=64)
        freqs = [0.1, 0.2]
        thetas = np.deg2rad([0, 90])
        stack, params = gabor_feature_stack(img, freqs, thetas)
        assert stack.shape == (4, 64, 64)
        assert len(params) == 4

    def test_feature_stack_nonnegative(self):
        img = _grating(n=64)
        stack, _ = gabor_feature_stack(img, [0.1], np.deg2rad([0, 90]))
        assert np.all(stack >= 0)


# ---------------------------------------------------------------------------
# Difference of Gaussians
# ---------------------------------------------------------------------------

class TestDoG:
    def test_blob_gives_positive_centre(self):
        img = np.zeros((64, 64))
        img[28:36, 28:36] = 1.0
        d = dog(img, 1.0, 2.0)
        # A bright blob on dark background => positive centre response.
        assert d[32, 32] > 0
        assert d[0, 0] == pytest.approx(0.0, abs=1e-6)

    def test_flat_image_gives_zero(self):
        img = np.full((32, 32), 0.5)
        d = dog(img, 1.0, 2.0)
        np.testing.assert_allclose(d, 0.0, atol=1e-9)

    def test_rejects_inverted_sigmas(self):
        with pytest.raises(ValueError):
            dog(np.zeros((8, 8)), 2.0, 1.0)