"""
Tests for BarcodeExtractor class.
"""

import pytest
import numpy as np
from medtda.barcodeextractor import BarcodeExtractor


class TestBarcodeExtractorInit:
    """Test BarcodeExtractor initialization."""
    
    def test_default_init(self):
        """Test initialization with default parameters."""
        extractor = BarcodeExtractor()
        
        assert extractor.preprocessor.spacing is None
        assert extractor.preprocessor.window is None
        assert extractor.preprocessor.normalize is False
        assert extractor.preprocessor.normalize_method == 'minmax'
        assert extractor.preprocessor.background_value is None
        assert extractor.filtration_type == 'sublevel'
        assert extractor.construction == 'T'
        assert extractor.max_dimension == -1
    
    def test_custom_init(self):
        """Test initialization with custom parameters."""
        extractor = BarcodeExtractor(
            spacing=(1.0, 1.0, 1.0),
            window=(50, 100),
            normalize=True,
            normalize_method='zscore',
            background_value=-1000,
            filtration_type='superlevel',
            construction='V',
            max_dimension=1
        )
        
        assert extractor.preprocessor.spacing == (1.0, 1.0, 1.0)
        assert extractor.preprocessor.window == (50, 100)
        assert extractor.preprocessor.normalize is True
        assert extractor.preprocessor.normalize_method == 'zscore'
        assert extractor.preprocessor.background_value == -1000
        assert extractor.filtration_type == 'superlevel'
        assert extractor.construction == 'V'
        assert extractor.max_dimension == 1
    
    def test_invalid_filtration_type(self):
        """Test that invalid filtration_type raises ValueError."""
        with pytest.raises(ValueError, match="filtration_type must be"):
            BarcodeExtractor(filtration_type='invalid')
    
    def test_invalid_construction(self):
        """Test that invalid construction raises ValueError."""
        with pytest.raises(ValueError, match="construction must be"):
            BarcodeExtractor(construction='invalid')
    
    def test_invalid_max_dimension(self):
        """Test that invalid max_dimension raises ValueError."""
        with pytest.raises(ValueError, match="max_dimension must be"):
            BarcodeExtractor(max_dimension='invalid')


class TestBarcodeExtractorSetters:
    """Test BarcodeExtractor setter methods."""
    
    def test_set_spacing(self):
        """Test set_spacing method."""
        extractor = BarcodeExtractor()
        extractor.set_spacing((1.0, 1.0, 1.0))
        assert extractor.preprocessor.spacing == (1.0, 1.0, 1.0)
    
    def test_set_window(self):
        """Test set_window method."""
        extractor = BarcodeExtractor()
        extractor.set_window(50, 100)
        assert extractor.preprocessor.window == (50, 100)
    
    def test_set_normalize(self):
        """Test set_normalize method."""
        extractor = BarcodeExtractor()
        extractor.set_normalize(True, 'zscore')
        assert extractor.preprocessor.normalize is True
        assert extractor.preprocessor.normalize_method == 'zscore'
    
    def test_set_background_value(self):
        """Test set_background_value method."""
        extractor = BarcodeExtractor()
        extractor.set_background_value(-1000)
        assert extractor.preprocessor.background_value == -1000
    
    def test_set_filtration_type(self):
        """Test set_filtration_type method."""
        extractor = BarcodeExtractor()
        extractor.set_filtration_type('superlevel')
        assert extractor.filtration_type == 'superlevel'
    
    def test_set_filtration_type_invalid(self):
        """Test set_filtration_type with invalid value."""
        extractor = BarcodeExtractor()
        with pytest.raises(ValueError, match="filtration_type must be"):
            extractor.set_filtration_type('invalid')
    
    def test_set_construction(self):
        """Test set_construction method."""
        extractor = BarcodeExtractor()
        extractor.set_construction('V')
        assert extractor.construction == 'V'
    
    def test_set_construction_invalid(self):
        """Test set_construction with invalid value."""
        extractor = BarcodeExtractor()
        with pytest.raises(ValueError, match="construction must be"):
            extractor.set_construction('invalid')
    
    def test_set_max_dimension(self):
        """Test set_max_dimension method."""
        extractor = BarcodeExtractor()
        extractor.set_max_dimension(1)
        assert extractor.max_dimension == 1
    
    def test_set_max_dimension_invalid(self):
        """Test set_max_dimension with invalid value."""
        extractor = BarcodeExtractor()
        with pytest.raises(ValueError, match="max_dimension must be"):
            extractor.set_max_dimension('invalid')


class TestBarcodeExtractorGetSettings:
    """Test BarcodeExtractor get_settings method."""
    
    def test_get_settings(self):
        """Test get_settings returns correct information."""
        extractor = BarcodeExtractor(
            normalize=True,
            normalize_method='minmax',
            filtration_type='superlevel',
            construction='V',
            max_dimension=1
        )
        
        settings = extractor.get_settings()
        
        assert 'preprocessing' in settings
        assert 'filtration_type' in settings
        assert 'construction' in settings
        assert 'max_dimension' in settings
        
        assert settings['filtration_type'] == 'superlevel'
        assert settings['construction'] == 'V'
        assert settings['max_dimension'] == 1
        assert settings['preprocessing']['normalize'] is True
        assert settings['preprocessing']['normalize_method'] == 'minmax'


class TestBarcodeExtractorExecute:
    """Test BarcodeExtractor execute method."""
    
    def test_execute_2d_array(self):
        """Test execute with 2D numpy array."""
        extractor = BarcodeExtractor(normalize=False)
        
        # Create simple 2D test image
        image = np.random.rand(20, 20)
        
        barcodes = extractor.execute(image)
        
        # Should have H0 and H1 for 2D image
        assert 'H0' in barcodes
        assert 'H1' in barcodes
        
        # Each barcode should be (n_features, 2)
        assert barcodes['H0'].ndim == 2
        assert barcodes['H0'].shape[1] == 2
        assert barcodes['H1'].ndim == 2
        assert barcodes['H1'].shape[1] == 2
    
    def test_execute_3d_array(self):
        """Test execute with 3D numpy array."""
        extractor = BarcodeExtractor(normalize=False)
        
        # Create simple 3D test image
        image = np.random.rand(10, 10, 10)
        
        barcodes = extractor.execute(image)
        
        # Should have H0, H1, H2 for 3D image
        assert 'H0' in barcodes
        assert 'H1' in barcodes
        assert 'H2' in barcodes
        
        # Each barcode should be (n_features, 2)
        assert barcodes['H0'].shape[1] == 2
        assert barcodes['H1'].shape[1] == 2
        assert barcodes['H2'].shape[1] == 2
    
    def test_execute_with_mask(self):
        """Test execute with mask."""
        extractor = BarcodeExtractor(
            normalize=True,
            normalize_method='minmax',
            background_value=-1
        )
        
        # Create test image and mask
        image = np.random.rand(20, 20)
        mask = np.ones((20, 20), dtype=bool)
        mask[:5, :5] = False  # Mask out top-left corner
        
        barcodes = extractor.execute(image, mask)
        
        assert 'H0' in barcodes
        assert 'H1' in barcodes
        assert barcodes['H0'].shape[1] == 2
        assert barcodes['H1'].shape[1] == 2
    
    def test_execute_with_normalization(self):
        """Test execute with normalization enabled."""
        extractor = BarcodeExtractor(
            normalize=True,
            normalize_method='minmax'
        )
        
        image = np.random.rand(20, 20) * 1000 + 500
        
        barcodes = extractor.execute(image)
        
        assert 'H0' in barcodes
        assert 'H1' in barcodes
        # Barcodes should exist
        assert barcodes['H0'].shape[0] >= 0
        assert barcodes['H1'].shape[0] >= 0
    
    def test_execute_with_specific_dimensions(self):
        """Test execute with specific max dimension."""
        extractor = BarcodeExtractor(
            normalize=False,
            max_dimension=0  # Only compute H0
        )
        
        image = np.random.rand(20, 20)
        
        barcodes = extractor.execute(image)
        
        # Should only have H0
        assert 'H0' in barcodes
        assert barcodes['H0'].shape[1] == 2
    
    def test_execute_superlevel_filtration(self):
        """Test execute with superlevel filtration."""
        extractor = BarcodeExtractor(
            normalize=False,
            filtration_type='superlevel'
        )
        
        image = np.random.rand(20, 20)
        
        barcodes = extractor.execute(image)
        
        assert 'H0' in barcodes
        assert 'H1' in barcodes
        assert barcodes['H0'].shape[1] == 2
        assert barcodes['H1'].shape[1] == 2
    
    def test_execute_v_construction(self):
        """Test execute with V-construction."""
        extractor = BarcodeExtractor(
            normalize=False,
            construction='V'
        )
        
        image = np.random.rand(20, 20)
        
        barcodes = extractor.execute(image)
        
        assert 'H0' in barcodes
        assert 'H1' in barcodes
        assert barcodes['H0'].shape[1] == 2
        assert barcodes['H1'].shape[1] == 2
    
    def test_execute_synthetic_circle(self):
        """Test execute on synthetic image with known topology (circle)."""
        extractor = BarcodeExtractor(
            normalize=False,
            filtration_type='superlevel'
        )
        
        # Create image with a circle (bright ring on dark background)
        image = np.zeros((30, 30))
        y, x = np.ogrid[-15:15, -15:15]
        circle_mask = (x**2 + y**2 >= 36) & (x**2 + y**2 <= 81)
        image[circle_mask] = 1.0
        
        barcodes = extractor.execute(image)
        
        # Should detect H1 features (the hole in the circle)
        assert 'H0' in barcodes
        assert 'H1' in barcodes
        assert barcodes['H1'].shape[0] > 0  # Should have at least one H1 feature
    
    def test_barcodes_birth_death_order(self):
        """Test that birth times are less than death times."""
        extractor = BarcodeExtractor(normalize=False)
        
        image = np.random.rand(20, 20)
        
        barcodes = extractor.execute(image)
        
        for dim_key, barcode in barcodes.items():
            if barcode.shape[0] > 0:  # If there are features
                births = barcode[:, 0]
                deaths = barcode[:, 1]
                # Birth should be <= death for all features
                assert np.all(births <= deaths), f"Invalid birth/death order in {dim_key}"

    def test_execute_with_multilabel_mask(self, sample_2d_multilabel_mask):
        """Test execute with multi-label mask and specific label."""
        extractor = BarcodeExtractor(
            normalize=True,
            background_value=-1,
            label=1  # Extract only label 1
        )
        
        image = np.ones((50, 50)) * 100
        barcodes = extractor.execute(image, sample_2d_multilabel_mask)
        
        assert 'H0' in barcodes
        assert 'H1' in barcodes
    
    def test_execute_override_label(self, sample_2d_multilabel_mask):
        """Test execute with label parameter override."""
        extractor = BarcodeExtractor(label=1)
        
        image = np.ones((50, 50)) * 100
        # Override instance label with label=2
        barcodes = extractor.execute(image, sample_2d_multilabel_mask, label=2)
        
        assert 'H0' in barcodes
        assert 'H1' in barcodes
    
    def test_set_label(self):
        """Test set_label method."""
        extractor = BarcodeExtractor()
        extractor.set_label(5)
        assert extractor.preprocessor.label == 5