"""
Tests for Preprocessor class.
"""

import pytest
import numpy as np
from pathlib import Path
import tempfile

from medtda.preprocessor import Preprocessor
from medtda.loaders import save_image


class TestPreprocessorInit:
    """Tests for Preprocessor initialization."""
    
    def test_default_init(self):
        """Test initialization with default parameters."""
        preprocessor = Preprocessor()
        assert preprocessor.spacing is None
        assert preprocessor.window is None
        assert preprocessor.normalize is False
        assert preprocessor.normalize_method == 'minmax'
        assert preprocessor.background_value is None
    
    def test_custom_init(self):
        """Test initialization with custom parameters."""
        preprocessor = Preprocessor(
            spacing=(1.0, 1.0, 1.0),
            window=(50, 100),
            normalize=True,
            normalize_method='zscore',
            background_value=-1.0
        )
        assert preprocessor.spacing == (1.0, 1.0, 1.0)
        assert preprocessor.window == (50, 100)
        assert preprocessor.normalize is True
        assert preprocessor.normalize_method == 'zscore'
        assert preprocessor.background_value == -1.0
    
    def test_invalid_normalize_method(self):
        """Test that invalid normalize_method raises ValueError."""
        with pytest.raises(ValueError, match="normalize_method must be"):
            Preprocessor(normalize_method='invalid')


class TestPreprocessorSetters:
    """Tests for Preprocessor setter methods."""
    
    def test_set_spacing(self):
        """Test set_spacing method."""
        preprocessor = Preprocessor()
        preprocessor.set_spacing((2.0, 2.0, 2.0))
        assert preprocessor.spacing == (2.0, 2.0, 2.0)
    
    def test_set_window(self):
        """Test set_window method."""
        preprocessor = Preprocessor()
        preprocessor.set_window(40, 80)
        assert preprocessor.window == (40, 80)
    
    def test_set_normalize(self):
        """Test set_normalize method."""
        preprocessor = Preprocessor()
        preprocessor.set_normalize(True, 'robust')
        assert preprocessor.normalize is True
        assert preprocessor.normalize_method == 'robust'
    
    def test_set_normalize_invalid_method(self):
        """Test that invalid method in set_normalize raises ValueError."""
        preprocessor = Preprocessor()
        with pytest.raises(ValueError, match="normalize_method must be"):
            preprocessor.set_normalize(True, 'invalid')
    
    def test_set_background_value(self):
        """Test set_background_value method."""
        preprocessor = Preprocessor()
        preprocessor.set_background_value(-500.0)
        assert preprocessor.background_value == -500.0


class TestPreprocessorPreprocess:
    """Tests for preprocess method."""
    
    def test_preprocess_array_no_mask(self):
        """Test preprocessing array without mask."""
        image = np.random.rand(50, 50) * 100
        preprocessor = Preprocessor()
        
        preprocessed, metadata = preprocessor.preprocess(image)
        
        assert isinstance(preprocessed, np.ndarray)
        assert preprocessed.shape == image.shape
        assert 'original_range' in metadata
        assert 'final_range' in metadata
        assert 'transforms' in metadata
    
    def test_preprocess_with_normalization(self):
        """Test preprocessing with normalization enabled."""
        image = np.random.rand(50, 50) * 100
        preprocessor = Preprocessor(normalize=True, normalize_method='minmax')
        
        preprocessed, metadata = preprocessor.preprocess(image)
        
        # Check that image is normalized to [0, 1]
        assert preprocessed.min() >= 0.0
        assert preprocessed.max() <= 1.0
        assert 'normalize_minmax' in metadata['transforms']
    
    def test_preprocess_with_mask(self):
        """Test preprocessing with mask."""
        image = np.random.rand(50, 50) * 100
        mask = np.zeros((50, 50), dtype=np.uint8)
        mask[10:40, 10:40] = 1
        
        preprocessor = Preprocessor(normalize=True, background_value=-1.0)
        preprocessed, metadata = preprocessor.preprocess(image, mask)
        
        # Check that background is set correctly
        assert preprocessed[0, 0] == -1.0  # Background pixel
        assert metadata['background_value'] == -1.0
        assert 'background_set' in metadata['transforms']
    
    def test_preprocess_auto_background_2d_normalized(self):
        """Test auto background value for 2D normalized image."""
        image = np.random.rand(30, 30) * 100
        mask = np.ones((30, 30), dtype=np.uint8)
        mask[0, 0] = 0
        
        preprocessor = Preprocessor(normalize=True)
        preprocessed, metadata = preprocessor.preprocess(image, mask)
        
        # Auto background for normalized should be -1
        assert metadata['background_value'] == -1.0
        assert preprocessed[0, 0] == -1.0
    
    def test_preprocess_auto_background_2d_unnormalized(self):
        """Test auto background value for 2D unnormalized image."""
        image = np.random.rand(30, 30) * 100
        mask = np.ones((30, 30), dtype=np.uint8)
        mask[0, 0] = 0
        
        preprocessor = Preprocessor(normalize=False)
        preprocessed, metadata = preprocessor.preprocess(image, mask)
        
        # Auto background for 2D unnormalized should be 0
        assert metadata['background_value'] == 0.0
        assert preprocessed[0, 0] == 0.0
    
    def test_preprocess_auto_background_3d_unnormalized(self):
        """Test auto background value for 3D unnormalized image."""
        image = np.random.rand(20, 20, 20) * 100
        mask = np.ones((20, 20, 20), dtype=np.uint8)
        mask[0, 0, 0] = 0
        
        preprocessor = Preprocessor(normalize=False)
        preprocessed, metadata = preprocessor.preprocess(image, mask)
        
        # Auto background for 3D unnormalized should be -3000
        assert metadata['background_value'] == -3000.0
        assert preprocessed[0, 0, 0] == -3000.0
    
    def test_preprocess_with_windowing(self):
        """Test preprocessing with windowing."""
        image = np.random.rand(50, 50) * 1000  # Values 0-1000
        preprocessor = Preprocessor(window=(100, 200))  # Window center=100, width=200
        
        preprocessed, metadata = preprocessor.preprocess(image)
        
        # Check that values are clipped to [0, 200]
        assert preprocessed.min() >= 0.0
        assert preprocessed.max() <= 200.0
        assert any('window' in t for t in metadata['transforms'])


class TestPreprocessorGetInfo:
    """Tests for get_preprocessing_info method."""
    
    def test_get_info(self):
        """Test that get_preprocessing_info returns correct settings."""
        preprocessor = Preprocessor(
            spacing=(1.0, 1.0, 1.0),
            window=(50, 100),
            normalize=True,
            normalize_method='zscore',
            background_value=-10.0
        )
        
        info = preprocessor.get_preprocessing_info()
        
        assert info['spacing'] == (1.0, 1.0, 1.0)
        assert info['window'] == (50, 100)
        assert info['normalize'] is True
        assert info['normalize_method'] == 'zscore'
        assert info['background_value'] == -10.0
        assert 'label' in info


class TestPreprocessorMultiLabelSupport:
    """Tests for multi-label mask support in Preprocessor."""
    
    def test_init_with_label(self):
        """Test initialization with label parameter."""
        preprocessor = Preprocessor(label=2)
        assert preprocessor.label == 2
    
    def test_set_label(self):
        """Test set_label method."""
        preprocessor = Preprocessor()
        preprocessor.set_label(3)
        assert preprocessor.label == 3
    
    def test_preprocess_with_single_label(self, sample_2d_multilabel_mask):
        """Test preprocessing with multi-label mask and specific label."""
        # Create non-uniform image with intensity gradient
        image = np.random.rand(50, 50) * 100 + 50  # Values between 50-150
        preprocessor = Preprocessor(normalize=True, background_value=-1.0, label=1, crop_to_roi=False)
        
        preprocessed, metadata = preprocessor.preprocess(image, sample_2d_multilabel_mask)
        
        # Check that only label 1 is preserved (values should not be background)
        assert preprocessed[15, 15] != -1.0  # Inside label 1 (normalized, but not background)
        assert preprocessed[30, 30] == -1.0  # Inside label 2 (should be background)
        assert preprocessed[0, 0] == -1.0  # Outside all labels
        assert metadata['label'] == 1
    
    def test_preprocess_with_different_label(self, sample_2d_multilabel_mask):
        """Test preprocessing with different label."""
        image = np.ones((50, 50)) * 100
        preprocessor = Preprocessor(label=2)
        
        preprocessed, metadata = preprocessor.preprocess(image, sample_2d_multilabel_mask, label=2)
        
        assert metadata['label'] == 2
    
    def test_preprocess_override_instance_label(self, sample_2d_multilabel_mask):
        """Test that execute-time label overrides instance label."""
        image = np.ones((50, 50)) * 100
        preprocessor = Preprocessor(label=1, background_value=0, crop_to_roi=False)
        
        # Override with label=2 at preprocess time
        preprocessed, metadata = preprocessor.preprocess(image, sample_2d_multilabel_mask, label=2)
        
        assert metadata['label'] == 2
        assert preprocessed[15, 15] == 0  # Label 1 area should be background
        assert preprocessed[30, 30] == 100  # Label 2 area should be preserved
    
    def test_preprocess_label_none_binary_mode(self, sample_2d_multilabel_mask):
        """Test preprocessing with label=None (binary mode)."""
        image = np.ones((50, 50)) * 100
        preprocessor = Preprocessor(background_value=0, crop_to_roi=False)
        
        preprocessed, metadata = preprocessor.preprocess(image, sample_2d_multilabel_mask, label=None)
        
        # All labels should be treated as foreground
        assert preprocessed[15, 15] == 100  # Label 1
        assert preprocessed[30, 30] == 100  # Label 2
        assert preprocessed[38, 15] == 100  # Label 3
        assert preprocessed[0, 0] == 0  # Background
        assert metadata['label'] is None


class TestPreprocessAndSave:
    """Tests for preprocess_and_save method."""
    
    @pytest.mark.skipif(
        not pytest.importorskip("PIL", reason="Pillow not installed"),
        reason="Pillow required for saving images"
    )
    def test_preprocess_and_save(self):
        """Test preprocess_and_save method."""
        with tempfile.TemporaryDirectory() as tmpdir:
            # Create and save original image
            image = np.random.rand(50, 50) * 100
            image_path = Path(tmpdir) / "original.png"
            save_image(image, image_path)
            
            # Preprocess and save
            preprocessor = Preprocessor(normalize=True)
            output_path = Path(tmpdir) / "preprocessed.png"
            
            preprocessed, metadata = preprocessor.preprocess_and_save(
                image_path,
                output_path=output_path
            )
            
            # Check that output file was created
            assert output_path.exists()
            assert isinstance(preprocessed, np.ndarray)
    
    @pytest.mark.skipif(
        not pytest.importorskip("PIL", reason="Pillow not installed"),
        reason="Pillow required for saving images"
    )
    def test_preprocess_and_save_auto_name(self):
        """Test preprocess_and_save with auto-generated output name."""
        with tempfile.TemporaryDirectory() as tmpdir:
            # Create and save original image
            image = np.random.rand(50, 50) * 100
            image_path = Path(tmpdir) / "original.png"
            save_image(image, image_path)
            
            # Preprocess and save with auto name
            preprocessor = Preprocessor(normalize=True)
            preprocessed, metadata = preprocessor.preprocess_and_save(image_path)
            
            # Check that default output file was created
            expected_path = Path(tmpdir) / "preprocessed_original.png"
            assert expected_path.exists()
