"""
Tests for FeatureExtractor class.
"""

import pytest
import numpy as np
from medtda.featureextractor import FeatureExtractor


class TestFeatureExtractorInit:
    """Test FeatureExtractor initialization."""
    
    def test_default_init(self):
        """Test initialization with default parameters."""
        extractor = FeatureExtractor()
        
        assert extractor.barcode_extractor.preprocessor.normalize is False
        assert extractor.barcode_extractor.filtration_type == 'sublevel'
        assert extractor.return_barcodes is False
        assert 'PersStats' in extractor.get_enabled_vectorization_methods()
    
    def test_custom_init(self):
        """Test initialization with custom parameters."""
        extractor = FeatureExtractor(
            normalize=True,
            normalize_method='zscore',
            filtration_type='superlevel',
            construction='V',
            max_dimension=1,
            return_barcodes=True,
            vectorization_method=['PersStats', 'BettiCurve']
        )
        
        assert extractor.barcode_extractor.preprocessor.normalize is True
        assert extractor.barcode_extractor.preprocessor.normalize_method == 'zscore'
        assert extractor.barcode_extractor.filtration_type == 'superlevel'
        assert extractor.barcode_extractor.construction == 'V'
        assert extractor.return_barcodes is True
        methods = extractor.get_enabled_vectorization_methods()
        assert 'PersStats' in methods
        assert 'BettiCurve' in methods
    
    def test_invalid_vectorization_method(self):
        """Test that invalid vectorization method raises ValueError."""
        with pytest.raises(ValueError, match="Unknown vectorization method"):
            FeatureExtractor(vectorization_method='InvalidMethod')


class TestFeatureExtractorSetters:
    """Test FeatureExtractor setter methods."""
    
    def test_set_spacing(self):
        """Test set_spacing method."""
        extractor = FeatureExtractor()
        extractor.set_spacing((1.0, 1.0, 1.0))
        assert extractor.barcode_extractor.preprocessor.spacing == (1.0, 1.0, 1.0)
    
    def test_set_window(self):
        """Test set_window method."""
        extractor = FeatureExtractor()
        extractor.set_window(50, 100)
        assert extractor.barcode_extractor.preprocessor.window == (50, 100)
    
    def test_set_normalize(self):
        """Test set_normalize method."""
        extractor = FeatureExtractor()
        extractor.set_normalize(True, 'zscore')
        assert extractor.barcode_extractor.preprocessor.normalize is True
        assert extractor.barcode_extractor.preprocessor.normalize_method == 'zscore'
    
    def test_set_background_value(self):
        """Test set_background_value method."""
        extractor = FeatureExtractor()
        extractor.set_background_value(-1000)
        assert extractor.barcode_extractor.preprocessor.background_value == -1000
    
    def test_set_filtration_type(self):
        """Test set_filtration_type method."""
        extractor = FeatureExtractor()
        extractor.set_filtration_type('superlevel')
        assert extractor.barcode_extractor.filtration_type == 'superlevel'
    
    def test_set_construction(self):
        """Test set_construction method."""
        extractor = FeatureExtractor()
        extractor.set_construction('V')
        assert extractor.barcode_extractor.construction == 'V'
    
    def test_set_max_dimension(self):
        """Test set_max_dimension method."""
        extractor = FeatureExtractor()
        extractor.set_max_dimension(1)
        assert extractor.barcode_extractor.max_dimension == 1
    
    def test_set_return_barcodes(self):
        """Test set_return_barcodes method."""
        extractor = FeatureExtractor()
        extractor.set_return_barcodes(True)
        assert extractor.return_barcodes is True


class TestVectorizationMethodManagement:
    """Test vectorization method management."""
    
    def test_set_vectorization_method(self):
        """Test setting a single vectorization method."""
        extractor = FeatureExtractor(vectorization_method='PersStats')
        extractor.set_vectorization_method('BettiCurve', resolution=50)
        
        methods = extractor.get_enabled_vectorization_methods()
        assert len(methods) == 1
        assert 'BettiCurve' in methods
        params = extractor.get_vectorization_params('BettiCurve')
        assert params['resolution'] == 50
    
    def test_set_vectorization_params(self):
        """Test updating parameters for a method."""
        extractor = FeatureExtractor(vectorization_method='BettiCurve')
        extractor.set_vectorization_params('BettiCurve', resolution=200)
        
        params = extractor.get_vectorization_params('BettiCurve')
        assert params['resolution'] == 200
    
    def test_set_vectorization_params_not_enabled(self):
        """Test setting params for non-enabled method raises error."""
        extractor = FeatureExtractor(vectorization_method='PersStats')
        
        with pytest.raises(ValueError, match="is not enabled"):
            extractor.set_vectorization_params('BettiCurve', resolution=50)
    
    def test_enable_vectorization_methods(self):
        """Test enabling multiple methods."""
        extractor = FeatureExtractor(vectorization_method='PersStats')
        extractor.enable_vectorization_methods(['BettiCurve', 'EntropySummary'])
        
        methods = extractor.get_enabled_vectorization_methods()
        assert 'PersStats' in methods
        assert 'BettiCurve' in methods
        assert 'EntropySummary' in methods
    
    def test_enable_all_vectorization_methods(self):
        """Test enabling all methods."""
        extractor = FeatureExtractor(vectorization_method='PersStats')
        extractor.enable_all_vectorization_methods()
        
        methods = extractor.get_enabled_vectorization_methods()
        # Should have all 8 methods
        assert len(methods) >= 8
        assert 'PersStats' in methods
        assert 'BettiCurve' in methods
        assert 'PersImage' in methods
    
    def test_disable_vectorization_method(self):
        """Test disabling a method."""
        extractor = FeatureExtractor(
            vectorization_method=['PersStats', 'BettiCurve']
        )
        extractor.disable_vectorization_method('BettiCurve')
        
        methods = extractor.get_enabled_vectorization_methods()
        assert 'PersStats' in methods
        assert 'BettiCurve' not in methods
    
    def test_get_enabled_vectorization_methods(self):
        """Test getting enabled methods."""
        extractor = FeatureExtractor(
            vectorization_method=['PersStats', 'BettiCurve']
        )
        methods = extractor.get_enabled_vectorization_methods()
        
        assert isinstance(methods, list)
        assert len(methods) == 2
        assert 'PersStats' in methods
        assert 'BettiCurve' in methods
    
    def test_get_vectorization_params(self):
        """Test getting method parameters."""
        extractor = FeatureExtractor(vectorization_method='PersStats')
        params = extractor.get_vectorization_params('PersStats')
        
        assert isinstance(params, dict)
    
    def test_get_vectorization_params_not_enabled(self):
        """Test getting params for non-enabled method returns empty dict."""
        extractor = FeatureExtractor(vectorization_method='PersStats')
        params = extractor.get_vectorization_params('BettiCurve')
        
        assert params == {}


class TestFeatureExtractorGetSettings:
    """Test get_settings method."""
    
    def test_get_settings(self):
        """Test get_settings returns complete information."""
        extractor = FeatureExtractor(
            normalize=True,
            filtration_type='superlevel',
            return_barcodes=True,
            vectorization_method=['PersStats', 'BettiCurve']
        )
        
        settings = extractor.get_settings()
        
        assert 'preprocessing' in settings
        assert 'filtration_type' in settings
        assert 'return_barcodes' in settings
        assert 'vectorization_methods' in settings
        assert 'vectorization_params' in settings
        
        assert settings['return_barcodes'] is True
        assert len(settings['vectorization_methods']) == 2


class TestFeatureExtractorExecute:
    """Test execute method."""
    
    def test_execute_2d_array_single_method(self):
        """Test execute with 2D array and single method."""
        extractor = FeatureExtractor(
            normalize=False,
            vectorization_method='PersStats'
        )
        
        image = np.random.rand(20, 20)
        features = extractor.execute(image)
        
        assert isinstance(features, dict)
        assert len(features) > 0
        
        # Check feature naming convention
        for feat_name in features.keys():
            assert feat_name.startswith('PersStats_H')
            assert '_' in feat_name
    
    def test_execute_with_return_barcodes(self):
        """Test execute with return_barcodes=True."""
        extractor = FeatureExtractor(
            normalize=False,
            vectorization_method='PersStats',
            return_barcodes=True
        )
        
        image = np.random.rand(20, 20)
        result = extractor.execute(image)
        
        assert isinstance(result, tuple)
        assert len(result) == 2
        
        features, barcodes = result
        assert isinstance(features, dict)
        assert isinstance(barcodes, dict)
        assert 'H0' in barcodes
        assert 'H1' in barcodes
    
    def test_execute_with_multiple_methods(self):
        """Test execute with multiple vectorization methods."""
        extractor = FeatureExtractor(
            normalize=False,
            vectorization_method=['PersStats', 'BettiCurve']
        )
        
        image = np.random.rand(20, 20)
        features = extractor.execute(image)
        
        # Should have features from both methods
        persstats_features = [k for k in features.keys() if 'PersStats' in k]
        betticurve_features = [k for k in features.keys() if 'BettiCurve' in k]
        
        assert len(persstats_features) > 0
        assert len(betticurve_features) > 0
    
    def test_execute_with_mask(self):
        """Test execute with mask."""
        extractor = FeatureExtractor(
            normalize=True,
            vectorization_method='PersStats'
        )
        
        image = np.random.rand(20, 20)
        mask = np.ones((20, 20), dtype=bool)
        mask[:5, :5] = False
        
        features = extractor.execute(image, mask)
        
        assert isinstance(features, dict)
        assert len(features) > 0
    
    def test_execute_3d_array(self):
        """Test execute with 3D array."""
        extractor = FeatureExtractor(
            normalize=False,
            vectorization_method='PersStats'
        )
        
        image = np.random.rand(10, 10, 10)
        features = extractor.execute(image)
        
        # Should have features for H0, H1, H2
        h0_features = [k for k in features.keys() if '_H0_' in k]
        h1_features = [k for k in features.keys() if '_H1_' in k]
        h2_features = [k for k in features.keys() if '_H2_' in k]
        
        assert len(h0_features) > 0
        assert len(h1_features) > 0
        assert len(h2_features) > 0
    
    def test_feature_naming_convention(self):
        """Test that feature names follow convention."""
        extractor = FeatureExtractor(
            normalize=False,
            vectorization_method='PersStats'
        )
        
        image = np.random.rand(15, 15)
        features = extractor.execute(image)
        
        for feat_name in features.keys():
            # Check format: {method}_{dimension}_{feature}
            parts = feat_name.split('_')
            assert parts[0] == 'PersStats'
            assert parts[1] in ['H0', 'H1']  # For 2D
            # parts[2:] is the feature name
    
    def test_execute_with_specific_dimensions(self):
        """Test execute with specific max dimension."""
        extractor = FeatureExtractor(
            normalize=False,
            vectorization_method='PersStats',
            max_dimension=0  # Only H0
        )
        
        image = np.random.rand(20, 20)
        features = extractor.execute(image)
        
        # Should only have H0 features
        h0_features = [k for k in features.keys() if '_H0_' in k]
        h1_features = [k for k in features.keys() if '_H1_' in k]
        
        assert len(h0_features) > 0
        assert len(h1_features) == 0
    
    def test_execute_superlevel_filtration(self):
        """Test execute with superlevel filtration."""
        extractor = FeatureExtractor(
            normalize=False,
            filtration_type='superlevel',
            vectorization_method='PersStats'
        )
        
        image = np.random.rand(20, 20)
        features = extractor.execute(image)
        
        assert isinstance(features, dict)
        assert len(features) > 0
    
    def test_features_are_numeric(self):
        """Test that all feature values are numeric."""
        extractor = FeatureExtractor(
            normalize=False,
            vectorization_method='PersStats'
        )
        
        image = np.random.rand(20, 20)
        features = extractor.execute(image)
        
        for feat_value in features.values():
            assert isinstance(feat_value, (int, float, np.number))
            assert not np.isnan(feat_value)
            assert not np.isinf(feat_value)

    def test_execute_with_multilabel_mask(self, sample_2d_multilabel_mask):
        """Test execute with multi-label mask and specific label."""
        extractor = FeatureExtractor(
            normalize=True,
            background_value=-1,
            label=1,  # Extract only label 1
            vectorization_method='PersStats'
        )
        
        image = np.ones((50, 50)) * 100
        features = extractor.execute(image, sample_2d_multilabel_mask)
        
        assert isinstance(features, dict)
        assert len(features) > 0
        # Check feature naming
        assert any('PersStats' in key for key in features.keys())
    
    def test_execute_override_label(self, sample_2d_multilabel_mask):
        """Test execute with label parameter override."""
        extractor = FeatureExtractor(label=1, vectorization_method='PersStats')
        
        image = np.ones((50, 50)) * 100
        # Override instance label with label=2
        features = extractor.execute(image, sample_2d_multilabel_mask, label=2)
        
        assert isinstance(features, dict)
        assert len(features) > 0
    
    def test_set_label(self):
        """Test set_label method."""
        extractor = FeatureExtractor()
        extractor.set_label(5)
        assert extractor.barcode_extractor.preprocessor.label == 5