"""
Tests for utils module.
"""

import pytest
import numpy as np
from medtda.utils import (
    normalize_image,
    apply_mask,
    check_array_dimensions,
    safe_divide,
    get_image_statistics,
    extract_label_from_mask,
    crop_to_roi
)


class TestExtractLabelFromMask:
    """Tests for extract_label_from_mask function."""
    
    def test_extract_single_label(self, sample_2d_multilabel_mask):
        """Test extracting a single label from multi-label mask."""
        binary_mask = extract_label_from_mask(sample_2d_multilabel_mask, label=1)
        assert binary_mask.dtype == np.uint8
        assert np.sum(binary_mask) == np.sum(sample_2d_multilabel_mask == 1)
        assert np.all(binary_mask[10:20, 10:20] == 1)
        assert np.all(binary_mask[25:35, 25:35] == 0)
    
    def test_extract_different_label(self, sample_2d_multilabel_mask):
        """Test extracting different labels."""
        binary_mask_2 = extract_label_from_mask(sample_2d_multilabel_mask, label=2)
        binary_mask_3 = extract_label_from_mask(sample_2d_multilabel_mask, label=3)
        assert np.sum(binary_mask_2) == np.sum(sample_2d_multilabel_mask == 2)
        assert np.sum(binary_mask_3) == np.sum(sample_2d_multilabel_mask == 3)
    
    def test_label_none_treats_as_binary(self, sample_2d_multilabel_mask):
        """Test that label=None treats mask as binary."""
        binary_mask = extract_label_from_mask(sample_2d_multilabel_mask, label=None)
        assert np.sum(binary_mask) == np.sum(sample_2d_multilabel_mask > 0)
        assert np.all(binary_mask[sample_2d_multilabel_mask > 0] == 1)
    
    def test_nonexistent_label_returns_empty(self, sample_2d_multilabel_mask):
        """Test that nonexistent label returns all zeros."""
        binary_mask = extract_label_from_mask(sample_2d_multilabel_mask, label=99)
        assert np.sum(binary_mask) == 0


class TestNormalizeImage:
    """Tests for normalize_image function."""

    def test_minmax_normalization(self):
        """Test min-max normalization."""
        image = np.array([[0, 5, 10], [15, 20, 25]])
        result = normalize_image(image, method='minmax')
        assert result.min() == 0.0
        assert result.max() == 1.0

    def test_zscore_normalization(self):
        """Test z-score normalization."""
        image = np.random.rand(50, 50) * 100
        result = normalize_image(image, method='zscore')
        assert np.abs(result.mean()) < 1e-10
        assert np.abs(result.std() - 1.0) < 1e-10

    def test_robust_normalization(self):
        """Test robust normalization."""
        image = np.random.rand(50, 50) * 100
        result = normalize_image(image, method='robust')
        assert isinstance(result, np.ndarray)
        assert result.shape == image.shape

    def test_with_mask(self):
        """Test normalization with mask."""
        image = np.random.rand(30, 30)
        mask = np.zeros((30, 30))
        mask[10:20, 10:20] = 1
        result = normalize_image(image, method='minmax', mask=mask)
        assert isinstance(result, np.ndarray)

    def test_invalid_method_raises_error(self):
        """Test that invalid method raises ValueError."""
        image = np.random.rand(10, 10)
        with pytest.raises(ValueError, match="Normalization method must be"):
            normalize_image(image, method='invalid')

    def test_mask_shape_mismatch_raises_error(self):
        """Test that mismatched mask shape raises ValueError."""
        image = np.random.rand(10, 10)
        mask = np.ones((5, 5))
        with pytest.raises(ValueError, match="Mask shape.*does not match"):
            normalize_image(image, mask=mask)
    
    def test_with_multilabel_mask(self, sample_2d_multilabel_mask):
        """Test normalization with multi-label mask and specific label."""
        image = np.random.rand(50, 50) * 100
        # Normalize using only label 2
        result = normalize_image(image, method='minmax', mask=sample_2d_multilabel_mask, label=2)
        assert isinstance(result, np.ndarray)
        assert result.shape == image.shape


class TestApplyMask:
    """Tests for apply_mask function."""

    def test_basic_masking(self):
        """Test basic mask application."""
        image = np.ones((10, 10)) * 100
        mask = np.zeros((10, 10))
        mask[3:7, 3:7] = 1
        result = apply_mask(image, mask, background_value=0)
        
        assert result[5, 5] == 100  # Inside mask
        assert result[0, 0] == 0    # Outside mask

    def test_custom_background_value(self):
        """Test mask with custom background value."""
        image = np.ones((10, 10)) * 100
        mask = np.zeros((10, 10))
        mask[5, 5] = 1
        result = apply_mask(image, mask, background_value=-999)
        
        assert result[5, 5] == 100
        assert result[0, 0] == -999

    def test_mask_shape_mismatch_raises_error(self):
        """Test that mismatched shapes raise ValueError."""
        image = np.ones((10, 10))
        mask = np.ones((5, 5))
        with pytest.raises(ValueError, match="Mask shape.*does not match"):
            apply_mask(image, mask)
    
    def test_with_multilabel_mask(self, sample_2d_multilabel_mask):
        """Test mask application with multi-label mask and specific label."""
        image = np.ones((50, 50)) * 100
        # Apply mask using only label 1
        result = apply_mask(image, sample_2d_multilabel_mask, background_value=0, label=1)
        assert result[15, 15] == 100  # Inside label 1
        assert result[30, 30] == 0    # Inside label 2 (but we're using label=1)
        assert result[0, 0] == 0      # Outside all labels
    
    def test_with_multilabel_mask_label_none(self, sample_2d_multilabel_mask):
        """Test mask application with label=None (binary mode)."""
        image = np.ones((50, 50)) * 100
        result = apply_mask(image, sample_2d_multilabel_mask, background_value=0, label=None)
        assert result[15, 15] == 100  # Inside label 1
        assert result[30, 30] == 100  # Inside label 2
        assert result[0, 0] == 0      # Outside all labels


class TestCheckArrayDimensions:
    """Tests for check_array_dimensions function."""

    def test_exact_dimensions(self):
        """Test checking for exact dimensions."""
        array = np.random.rand(10, 10)
        check_array_dimensions(array, expected_dims=2)  # Should not raise

    def test_multiple_valid_dimensions(self):
        """Test checking for multiple valid dimensions."""
        array = np.random.rand(10, 10, 10)
        check_array_dimensions(array, expected_dims=(2, 3, 4))  # Should not raise

    def test_invalid_dimensions_raises_error(self):
        """Test that invalid dimensions raise ValueError."""
        array = np.random.rand(10, 10, 10)
        with pytest.raises(ValueError, match="must have.*dimensions"):
            check_array_dimensions(array, expected_dims=2)

    def test_custom_name_in_error(self):
        """Test that custom name appears in error message."""
        array = np.random.rand(10)
        with pytest.raises(ValueError, match="my_custom_array"):
            check_array_dimensions(array, expected_dims=2, name="my_custom_array")


class TestSafeDivide:
    """Tests for safe_divide function."""

    def test_normal_division(self):
        """Test normal division."""
        numerator = np.array([10, 20, 30])
        denominator = np.array([2, 4, 5])
        result = safe_divide(numerator, denominator)
        expected = np.array([5, 5, 6])
        np.testing.assert_array_equal(result, expected)

    def test_division_by_zero(self):
        """Test division by zero returns zero."""
        numerator = np.array([10, 20, 30])
        denominator = np.array([2, 0, 5])
        result = safe_divide(numerator, denominator)
        assert result[0] == 5
        assert result[1] == 0  # Division by zero
        assert result[2] == 6

    def test_handles_inf(self):
        """Test that inf values are replaced with zero."""
        numerator = np.array([np.inf, 10, -np.inf])
        denominator = np.array([1, 2, 1])
        result = safe_divide(numerator, denominator)
        assert result[0] == 0
        assert result[1] == 5
        assert result[2] == 0


class TestGetImageStatistics:
    """Tests for get_image_statistics function."""

    def test_basic_statistics(self):
        """Test computation of basic statistics."""
        image = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
        stats = get_image_statistics(image)
        
        assert stats['min'] == 1
        assert stats['max'] == 10
        assert stats['mean'] == 5.5
        assert stats['median'] == 5.5

    def test_with_mask(self):
        """Test statistics with mask."""
        image = np.arange(100).reshape(10, 10)
        mask = np.zeros((10, 10))
        mask[0:5, 0:5] = 1
        
        stats = get_image_statistics(image, mask=mask)
        assert isinstance(stats, dict)
        assert all(key in stats for key in ['min', 'max', 'mean', 'std', 'median', 'q1', 'q3'])

    def test_empty_mask(self):
        """Test statistics with empty mask."""
        image = np.random.rand(10, 10)
        mask = np.zeros((10, 10))
        
        stats = get_image_statistics(image, mask=mask)
        assert all(stats[key] == 0 for key in stats.keys())

    def test_2d_image(self):
        """Test statistics for 2D image."""
        image = np.random.rand(50, 50)
        stats = get_image_statistics(image)
        
        assert 0 <= stats['min'] <= 1
        assert 0 <= stats['max'] <= 1
        assert stats['q1'] <= stats['median'] <= stats['q3']


class TestCropToRoi:
    """Tests for crop_to_roi function."""
    
    def test_basic_cropping_2d(self):
        """Test basic ROI cropping in 2D."""
        image = np.random.rand(100, 100)
        mask = np.zeros((100, 100))
        mask[40:60, 40:60] = 1  # Small ROI
        
        cropped_img, cropped_mask, info = crop_to_roi(image, mask, padding=1)
        
        # Should be 20x20 + 2 padding on each side = 22x22
        assert cropped_img.shape == (22, 22)
        assert cropped_mask.shape == (22, 22)
        assert info['original_shape'] == (100, 100)
        assert info['cropped_shape'] == (22, 22)
        assert info['roi_size'] == 400  # 20x20 pixels
        assert info['reduction_factor'] > 1
    
    def test_basic_cropping_3d(self):
        """Test basic ROI cropping in 3D."""
        image = np.random.rand(100, 100, 100)
        mask = np.zeros((100, 100, 100))
        mask[40:60, 40:60, 40:60] = 1  # Cubic ROI
        
        cropped_img, cropped_mask, info = crop_to_roi(image, mask, padding=1)
        
        # Should be 20x20x20 + 2 padding on each side = 22x22x22
        assert cropped_img.shape == (22, 22, 22)
        assert cropped_mask.shape == (22, 22, 22)
        assert info['roi_size'] == 8000  # 20x20x20 voxels
    
    def test_custom_padding(self):
        """Test custom padding value."""
        image = np.random.rand(50, 50)
        mask = np.zeros((50, 50))
        mask[20:30, 20:30] = 1
        
        # No padding
        cropped_img_0, _, info_0 = crop_to_roi(image, mask, padding=0)
        assert cropped_img_0.shape == (10, 10)
        
        # 5 pixel padding
        cropped_img_5, _, info_5 = crop_to_roi(image, mask, padding=5)
        assert cropped_img_5.shape == (20, 20)  # 10 + 5*2
    
    def test_edge_roi(self):
        """Test ROI at image edge."""
        image = np.random.rand(50, 50)
        mask = np.zeros((50, 50))
        mask[0:10, 0:10] = 1  # ROI at corner
        
        cropped_img, cropped_mask, info = crop_to_roi(image, mask, padding=5)
        
        # Padding should be clipped at edges
        assert cropped_img.shape[0] <= 15  # 10 + 5, but clipped at edge
        assert cropped_img.shape[1] <= 15
        assert np.all(cropped_img.shape == cropped_mask.shape)
    
    def test_multilabel_mask(self):
        """Test with multi-label mask."""
        image = np.random.rand(50, 50)
        mask = np.zeros((50, 50))
        mask[10:20, 10:20] = 1
        mask[25:35, 25:35] = 2
        mask[40:45, 40:45] = 3
        
        # Extract label 2
        cropped_img, cropped_mask, info = crop_to_roi(image, mask, padding=1, label=2)
        
        # Should only crop to label 2's ROI
        assert cropped_img.shape == (12, 12)  # 10 + 2 padding
        assert np.sum(cropped_mask == 2) == 100  # 10x10 pixels with label 2
    
    def test_label_none_uses_all_nonzero(self):
        """Test that label=None uses all non-zero mask values."""
        image = np.random.rand(60, 60)
        mask = np.zeros((60, 60))
        mask[10:20, 10:20] = 1
        mask[40:50, 40:50] = 2
        
        cropped_img, cropped_mask, info = crop_to_roi(image, mask, padding=0, label=None)
        
        # Should crop to bounding box containing both labels
        assert cropped_img.shape == (40, 40)  # from (10, 10) to (50, 50)
    
    def test_empty_mask_raises_error(self):
        """Test that empty mask raises ValueError."""
        image = np.random.rand(50, 50)
        mask = np.zeros((50, 50))
        
        with pytest.raises(ValueError, match="Mask is all zeros"):
            crop_to_roi(image, mask, padding=1)
    
    def test_shape_mismatch_raises_error(self):
        """Test that shape mismatch raises ValueError."""
        image = np.random.rand(50, 50)
        mask = np.zeros((40, 40))
        
        with pytest.raises(ValueError, match="does not match"):
            crop_to_roi(image, mask, padding=1)
    
    def test_reduction_factor(self):
        """Test reduction factor calculation."""
        image = np.random.rand(100, 100, 100)
        mask = np.zeros((100, 100, 100))
        mask[45:55, 45:55, 45:55] = 1  # Very small ROI
        
        cropped_img, cropped_mask, info = crop_to_roi(image, mask, padding=1)
        
        original_size = 100 * 100 * 100
        cropped_size = np.prod(cropped_img.shape)
        expected_reduction = original_size / cropped_size
        
        assert abs(info['reduction_factor'] - expected_reduction) < 0.01
        assert info['reduction_factor'] > 10  # Should be much smaller
    
    def test_bbox_slice_format(self):
        """Test that bbox is returned as tuple of slices."""
        image = np.random.rand(50, 50)
        mask = np.zeros((50, 50))
        mask[20:30, 20:30] = 1
        
        _, _, info = crop_to_roi(image, mask, padding=1)
        
        assert 'bbox' in info
        assert isinstance(info['bbox'], tuple)
        assert all(isinstance(s, slice) for s in info['bbox'])
        
        # Verify we can use bbox to crop
        manual_crop = image[info['bbox']]
        auto_crop, _, _ = crop_to_roi(image, mask, padding=1)
        assert np.array_equal(manual_crop, auto_crop)
    
    def test_preserves_image_values(self):
        """Test that cropping preserves original image values."""
        # Create image with known pattern
        image = np.arange(2500).reshape(50, 50)
        mask = np.zeros((50, 50))
        mask[20:30, 20:30] = 1
        
        cropped_img, _, info = crop_to_roi(image, mask, padding=0)
        
        # Verify values are preserved
        assert np.array_equal(cropped_img, image[20:30, 20:30])

