"""
Tests for persistent homology computation module.
"""

import pytest
import numpy as np
from medtda.ph_computer import compute_ph, clean_and_cap_ph


class TestCleanAndCapPH:
    """Tests for clean_and_cap_ph function."""

    def test_removes_zero_persistence(self):
        """Test that zero-persistence bars are removed."""
        ph = np.array([
            [0, 0.0, 0.0],  # Zero persistence
            [0, 0.0, 1.0],  # Non-zero persistence
            [1, 2.0, 2.0],  # Zero persistence
            [1, 2.0, 5.0],  # Non-zero persistence
        ])
        result = clean_and_cap_ph(ph, cap_value=10.0)
        assert result.shape[0] == 2
        assert np.all(result[:, 2] - result[:, 1] > 0)

    def test_caps_infinity(self):
        """Test that infinite values are capped."""
        ph = np.array([
            [0, 0.0, np.inf],
            [1, 2.0, 1e400],
        ])
        result = clean_and_cap_ph(ph, cap_value=10.0)
        assert np.all(result[:, 2] == 10.0)

    def test_preserves_finite_values(self):
        """Test that finite values are preserved."""
        ph = np.array([
            [0, 0.0, 5.0],
            [1, 2.0, 8.0],
        ])
        result = clean_and_cap_ph(ph, cap_value=10.0)
        np.testing.assert_array_equal(result, ph)


class TestComputePH:
    """Tests for compute_ph function."""

    def test_2d_image_sublevel(self):
        """Test PH computation on 2D image with sublevel filtration."""
        image = np.random.rand(20, 20).astype(np.float32)
        ph_array, dgms = compute_ph(image, maxdim=1, construction='T', filtration='sublevel')
        
        assert isinstance(ph_array, np.ndarray)
        assert isinstance(dgms, list)
        assert len(dgms) == 2  # H0 and H1

    def test_3d_image_sublevel(self):
        """Test PH computation on 3D image with sublevel filtration."""
        image = np.random.rand(10, 10, 10).astype(np.float32)
        ph_array, dgms = compute_ph(image, maxdim=2, construction='T', filtration='sublevel')
        
        assert isinstance(ph_array, np.ndarray)
        assert isinstance(dgms, list)
        assert len(dgms) == 3  # H0, H1, and H2

    def test_superlevel_filtration(self):
        """Test PH computation with superlevel filtration."""
        image = np.random.rand(15, 15).astype(np.float32)
        ph_array, dgms = compute_ph(image, maxdim=1, construction='T', filtration='superlevel')
        
        assert isinstance(ph_array, np.ndarray)
        assert isinstance(dgms, list)

    def test_v_construction(self):
        """Test PH computation with V-construction."""
        image = np.random.rand(15, 15).astype(np.float32)
        ph_array, dgms = compute_ph(image, maxdim=1, construction='V', filtration='sublevel')
        
        assert isinstance(ph_array, np.ndarray)
        assert isinstance(dgms, list)

    def test_invalid_construction_raises_error(self):
        """Test that invalid construction type raises ValueError."""
        image = np.random.rand(10, 10)
        with pytest.raises(ValueError, match="Construction must be"):
            compute_ph(image, construction='invalid')

    def test_invalid_filtration_raises_error(self):
        """Test that invalid filtration type raises ValueError."""
        image = np.random.rand(10, 10)
        with pytest.raises(ValueError, match="Filtration must be"):
            compute_ph(image, filtration='invalid')

    def test_max_value_capping(self):
        """Test that max_value parameter caps infinite persistence."""
        image = np.random.rand(15, 15).astype(np.float32)
        ph_array, dgms = compute_ph(image, maxdim=1, max_value=100.0)
        
        # Check that all death values are <= max_value
        assert np.all(ph_array[:, 2] <= 100.0)

    def test_auto_maxdim_2d(self):
        """Test that maxdim is auto-calculated as 1 for 2D images."""
        image = np.random.rand(20, 20).astype(np.float32)
        ph_array, dgms = compute_ph(image)  # maxdim=None (default)
        
        assert isinstance(ph_array, np.ndarray)
        assert isinstance(dgms, list)
        assert len(dgms) == 2  # H0 and H1

    def test_auto_maxdim_3d(self):
        """Test that maxdim is auto-calculated as 2 for 3D images."""
        image = np.random.rand(10, 10, 10).astype(np.float32)
        ph_array, dgms = compute_ph(image)  # maxdim=None (default)
        
        assert isinstance(ph_array, np.ndarray)
        assert isinstance(dgms, list)
        assert len(dgms) == 3  # H0, H1, and H2

    def test_auto_maxdim_4d(self):
        """Test that maxdim is auto-calculated as 3 for 4D images."""
        image = np.random.rand(8, 8, 8, 5).astype(np.float32)
        ph_array, dgms = compute_ph(image)  # maxdim=None (default)
        
        assert isinstance(ph_array, np.ndarray)
        assert isinstance(dgms, list)
        assert len(dgms) == 4  # H0, H1, H2, and H3
