"""
Test suite for Sen4MapDatasetMonthlyComposites dataset.

This module tests the sen4map.py dataset implementation, covering:
- Initialization with various parameters
- HDF5 data loading and processing
- Monthly composite computation via median
- SCL masking
- Band selection and channel mapping
- Classification map selection (land-cover vs crops)
- Preprocessing operations (crop, resize, reverse tiling)
- Min-max normalization
- Edge cases and error handling
"""

import pytest
import numpy as np
import torch
import h5py
import pickle
from unittest.mock import Mock, patch, MagicMock, mock_open
from torchvision.transforms.v2 import InterpolationMode

from terratorch.datasets.sen4map import Sen4MapDatasetMonthlyComposites
from terratorch.datasets.utils import HLSBands


@pytest.fixture
def mock_h5_data():
    """Create a mock HDF5 file structure with sample data."""
    # Create mock HDF5 file
    mock_h5 = MagicMock(spec=h5py.File)
    
    # Create sample data for one patch
    # Shape: (temporal, height, width) for each band
    n_temporal = 12  # 12 acquisitions across the year
    height, width = 120, 120
    
    # Create numpy arrays for each band (not mocks, real data)
    band_data = {}
    band_names = ['B2', 'B3', 'B4', 'B5', 'B6', 'B7', 'B8', 'B8A', 'B11', 'B12', 'SCL']
    
    for band_name in band_names:
        if band_name == 'SCL':
            # SCL mask with values < 9 (good quality)
            band_data[band_name] = np.random.randint(0, 9, size=(n_temporal, height, width))
        else:
            # Spectral bands with realistic values
            band_data[band_name] = np.random.randint(100, 3000, size=(n_temporal, height, width)).astype(np.float32)
    
    # Create dictionary-like mock that returns numpy arrays
    def getitem_patch(key):
        return band_data[key]
    
    mock_patch = MagicMock()
    mock_patch.__getitem__.side_effect = getitem_patch
    
    # Create mock attributes with Image_ID for temporal indexing
    # Distribute images across 12 months of 2018
    image_ids = np.array([f'S2_2018{month:02d}15_patch_001' for month in range(1, 13)])
    
    # Mock attrs that behaves like a dict
    mock_attrs = MagicMock()
    mock_attrs.__getitem__ = lambda self, key: image_ids if key == 'Image_ID' else 'A10'
    mock_patch.attrs = mock_attrs
    
    # Set up the mock h5 file to return the patch
    mock_h5.keys.return_value = ['patch_0001', 'patch_0002', 'patch_0003']
    mock_h5.__getitem__.return_value = mock_patch
    
    return mock_h5


@pytest.fixture
def mock_h5_crops_data():
    """Create a mock HDF5 file structure for crops classification."""
    mock_h5 = MagicMock(spec=h5py.File)
    
    n_temporal = 12
    height, width = 120, 120
    
    band_data = {}
    band_names = ['B2', 'B3', 'B4', 'B5', 'B6', 'B7', 'B8', 'B8A', 'B11', 'B12', 'SCL']
    
    for band_name in band_names:
        if band_name == 'SCL':
            band_data[band_name] = np.random.randint(0, 9, size=(n_temporal, height, width))
        else:
            band_data[band_name] = np.random.randint(100, 3000, size=(n_temporal, height, width)).astype(np.float32)
    
    def getitem_patch(key):
        return band_data[key]
    
    mock_patch = MagicMock()
    mock_patch.__getitem__.side_effect = getitem_patch
    
    # Use a crop label
    image_ids = np.array([f'S2_2018{month:02d}15_patch_001' for month in range(1, 13)])
    mock_attrs = MagicMock()
    mock_attrs.__getitem__ = lambda self, key: image_ids if key == 'Image_ID' else 'B11'
    mock_patch.attrs = mock_attrs
    
    mock_h5.keys.return_value = ['patch_0001']
    mock_h5.__getitem__.return_value = mock_patch
    
    return mock_h5


class TestSen4MapDatasetInitialization:
    """Test dataset initialization and parameter validation."""
    
    def test_basic_initialization(self, mock_h5_data):
        """Test basic dataset initialization with minimal parameters."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        assert dataset.h5data == mock_h5_data
        assert dataset.h5data_keys == ['patch_0001', 'patch_0002', 'patch_0003']
        assert dataset.crop_size is None
        assert dataset.input_channels is None
        assert dataset.classification_map == Sen4MapDatasetMonthlyComposites.land_cover_classification_map
        assert dataset.resize is False
        assert dataset.reverse_tile is False
    
    def test_initialization_with_crop_size(self, mock_h5_data):
        """Test initialization with crop size parameter."""
        crop_size = 64
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            crop_size=crop_size
        )
        
        assert dataset.crop_size == crop_size
    
    def test_initialization_with_custom_keys(self, mock_h5_data):
        """Test initialization with custom h5data_keys."""
        custom_keys = ['patch_0001', 'patch_0003']
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            h5data_keys=custom_keys
        )
        
        assert dataset.h5data_keys == custom_keys
        assert len(dataset) == 2
    
    def test_initialization_with_band_selection(self, mock_h5_data):
        """Test initialization with band selection."""
        dataset_bands = [HLSBands.BLUE, HLSBands.GREEN, HLSBands.RED, HLSBands.NIR_NARROW]
        input_bands = [HLSBands.RED, HLSBands.GREEN, HLSBands.BLUE]
        
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            dataset_bands=dataset_bands,
            input_bands=input_bands
        )
        
        # RED=2, GREEN=1, BLUE=0 in dataset_bands
        expected_channels = [2, 1, 0]
        assert dataset.input_channels == expected_channels
    
    def test_initialization_input_bands_without_dataset_bands_raises_error(self, mock_h5_data):
        """Test that providing input_bands without dataset_bands raises ValueError."""
        with pytest.raises(ValueError, match="input_bands was provided without specifying the dataset_bands"):
            Sen4MapDatasetMonthlyComposites(
                h5py_file_object=mock_h5_data,
                input_bands=[HLSBands.RED, HLSBands.GREEN]
            )
    
    def test_initialization_with_resize_parameters(self, mock_h5_data):
        """Test initialization with resize parameters."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            resize=True,
            resize_to=[256, 256],
            resize_interpolation=InterpolationMode.NEAREST,
            resize_antialiasing=False
        )
        
        assert dataset.resize is True
        assert dataset.resize_to == [256, 256]
        assert dataset.resize_interpolation == InterpolationMode.NEAREST
        assert dataset.resize_antialiasing is False
    
    def test_initialization_with_reverse_tile(self, mock_h5_data):
        """Test initialization with reverse tiling."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            reverse_tile=True,
            reverse_tile_size=5
        )
        
        assert dataset.reverse_tile is True
        assert dataset.reverse_tile_size == 5
    
    def test_initialization_with_crops_classification(self, mock_h5_crops_data):
        """Test initialization with crops classification map."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_crops_data,
            classification_map="crops",
            h5data_keys=['patch_0001']
        )
        
        assert dataset.classification_map == Sen4MapDatasetMonthlyComposites.crop_classification_map
    
    def test_initialization_invalid_classification_map_raises_error(self, mock_h5_data):
        """Test that invalid classification_map raises ValueError."""
        with pytest.raises(ValueError, match="is not from the list of valid ones"):
            Sen4MapDatasetMonthlyComposites(
                h5py_file_object=mock_h5_data,
                classification_map="invalid-map"
            )
    
    @patch("builtins.open", new_callable=mock_open)
    @patch("pickle.dump")
    def test_initialization_saves_keys(self, mock_pickle_dump, mock_file, mock_h5_data):
        """Test that dataset saves keys when save_keys_path is provided."""
        save_path = "/tmp/keys.pkl"
        
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            save_keys_path=save_path
        )
        
        mock_file.assert_called_once_with(save_path, "wb")
        mock_pickle_dump.assert_called_once()
        # Verify the keys were saved
        saved_keys = mock_pickle_dump.call_args[0][0]
        assert saved_keys == ['patch_0001', 'patch_0002', 'patch_0003']
    
    def test_initialization_crops_without_keys_prints_warning(self, mock_h5_crops_data, capsys):
        """Test that crops classification without keys prints a warning."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_crops_data,
            classification_map="crops"
        )
        
        captured = capsys.readouterr()
        assert "Crop classification task chosen but no keys supplied" in captured.out


class TestSen4MapDatasetGetItem:
    """Test __getitem__ method and data loading."""
    
    def test_getitem_basic(self, mock_h5_data):
        """Test basic __getitem__ functionality."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        sample = dataset[0]
        
        assert "image" in sample
        assert "label" in sample
        assert isinstance(sample["image"], torch.Tensor)
        assert isinstance(sample["label"], torch.LongTensor)
        
        # Check image shape: (channels=10, temporal=12, height, width)
        image = sample["image"]
        assert image.shape[0] == 10  # 10 spectral bands
        assert image.shape[1] == 12  # 12 monthly composites
        assert image.ndim == 4
    
    def test_getitem_with_crop(self, mock_h5_data):
        """Test __getitem__ with center cropping."""
        crop_size = 64
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            crop_size=crop_size
        )
        
        sample = dataset[0]
        image = sample["image"]
        
        # Check that spatial dimensions are cropped
        assert image.shape[2] == crop_size
        assert image.shape[3] == crop_size
    
    def test_getitem_with_resize(self, mock_h5_data):
        """Test __getitem__ with resizing."""
        resize_to = [224, 224]
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            resize=True,
            resize_to=resize_to
        )
        
        sample = dataset[0]
        image = sample["image"]
        
        # Check that image is resized
        assert image.shape[2] == resize_to[0]
        assert image.shape[3] == resize_to[1]
    
    def test_getitem_with_reverse_tile(self, mock_h5_data):
        """Test __getitem__ with reverse tiling."""
        crop_size = 64
        tile_size = 3
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            crop_size=crop_size,
            reverse_tile=True,
            reverse_tile_size=tile_size
        )
        
        sample = dataset[0]
        image = sample["image"]
        
        # After reverse tiling, spatial dimensions should be multiplied by tile_size
        expected_size = crop_size * tile_size
        assert image.shape[2] == expected_size
        assert image.shape[3] == expected_size
    
    def test_getitem_with_band_selection(self, mock_h5_data):
        """Test __getitem__ with band selection."""
        dataset_bands = [HLSBands.BLUE, HLSBands.GREEN, HLSBands.RED, 
                        HLSBands.RED_EDGE_1, HLSBands.RED_EDGE_2, HLSBands.RED_EDGE_3,
                        HLSBands.NIR_BROAD, HLSBands.NIR_NARROW, HLSBands.SWIR_1, HLSBands.SWIR_2]
        input_bands = [HLSBands.RED, HLSBands.GREEN, HLSBands.BLUE]
        
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            dataset_bands=dataset_bands,
            input_bands=input_bands
        )
        
        sample = dataset[0]
        image = sample["image"]
        
        # Should have 3 selected channels
        assert image.shape[0] == 3
    
    def test_getitem_normalization_clips_to_01(self, mock_h5_data):
        """Test that normalization clips values to [0, 1]."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        sample = dataset[0]
        image = sample["image"]
        
        # Check that values are clipped to [0, 1]
        assert image.min() >= 0.0
        assert image.max() <= 1.0
    
    def test_getitem_label_mapping_land_cover(self, mock_h5_data):
        """Test label mapping for land-cover classification."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            classification_map="land-cover"
        )
        
        sample = dataset[0]
        label = sample["label"]
        
        # 'A10' should map to class 0 in land_cover_classification_map
        assert label.item() == 0
    
    def test_getitem_label_mapping_crops(self, mock_h5_crops_data):
        """Test label mapping for crops classification."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_crops_data,
            classification_map="crops",
            h5data_keys=['patch_0001']
        )
        
        sample = dataset[0]
        label = sample["label"]
        
        # 'B11' should map to class 0 (Cereals) in crop_classification_map
        assert label.item() == 0


class TestSen4MapDatasetGetData:
    """Test get_data method and monthly composite computation."""
    
    def test_get_data_computes_monthly_composites(self, mock_h5_data):
        """Test that get_data computes monthly composites correctly."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        im = mock_h5_data['patch_0001']
        image, label = dataset.get_data(im)
        
        # Should have 12 monthly composites
        assert image.shape[1] == 12
        # Should have 10 spectral bands
        assert image.shape[0] == 10
        assert isinstance(image, torch.Tensor)
    
    def test_get_data_applies_scl_mask(self, mock_h5_data):
        """Test that SCL mask is applied correctly."""
        # Create mock data with SCL values >= 9 (bad quality)
        mock_h5 = MagicMock(spec=h5py.File)
        n_temporal, height, width = 12, 120, 120
        
        # Create real numpy arrays
        band_data = {}
        # Set SCL to values >= 9
        band_data['SCL'] = np.full((n_temporal, height, width), 10)
        
        # Set band data to non-zero values
        for band in ['B2', 'B3', 'B4', 'B5', 'B6', 'B7', 'B8', 'B8A', 'B11', 'B12']:
            band_data[band] = np.full((n_temporal, height, width), 1000.0)
        
        def getitem_patch(key):
            return band_data[key]
        
        mock_patch = MagicMock()
        mock_patch.__getitem__.side_effect = getitem_patch
        
        image_ids = np.array([f'S2_2018{month:02d}15_patch_001' for month in range(1, 13)])
        mock_attrs = MagicMock()
        mock_attrs.__getitem__ = lambda self, key: image_ids if key == 'Image_ID' else 'A10'
        mock_patch.attrs = mock_attrs
        
        mock_h5.keys.return_value = ['patch_0001']
        mock_h5.__getitem__.return_value = mock_patch
        
        dataset = Sen4MapDatasetMonthlyComposites(h5py_file_object=mock_h5)
        image, label = dataset.get_data(mock_patch)
        
        # With SCL >= 9, all values should be masked to 0
        # After normalization and clipping, values should be near 0
        assert torch.all(image == 0)
    
    def test_get_data_with_crop_size(self, mock_h5_data):
        """Test get_data with center cropping."""
        crop_size = 64
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            crop_size=crop_size
        )
        
        im = mock_h5_data['patch_0001']
        image, label = dataset.get_data(im)
        
        assert image.shape[2] == crop_size
        assert image.shape[3] == crop_size
    
    def test_get_data_with_resize(self, mock_h5_data):
        """Test get_data with resizing."""
        resize_to = [224, 224]
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            resize=True,
            resize_to=resize_to
        )
        
        im = mock_h5_data['patch_0001']
        image, label = dataset.get_data(im)
        
        assert image.shape[2] == resize_to[0]
        assert image.shape[3] == resize_to[1]
    
    def test_get_data_with_reverse_tile(self, mock_h5_data):
        """Test get_data with reverse tiling."""
        crop_size = 64
        tile_size = 3
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            crop_size=crop_size,
            reverse_tile=True,
            reverse_tile_size=tile_size
        )
        
        im = mock_h5_data['patch_0001']
        image, label = dataset.get_data(im)
        
        expected_size = crop_size * tile_size
        assert image.shape[2] == expected_size
        assert image.shape[3] == expected_size
    
    def test_get_data_label_extraction(self, mock_h5_data):
        """Test label extraction from HDF5 attributes."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        im = mock_h5_data['patch_0001']
        image, label = dataset.get_data(im)
        
        # Label should be extracted and mapped
        assert isinstance(label, np.ndarray)
        assert label.dtype == np.float32


class TestSen4MapDatasetCropCenter:
    """Test crop_center method."""
    
    def test_crop_center_basic(self, mock_h5_data):
        """Test basic center cropping."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        # Create a test tensor (channels, temporal, height, width)
        img = torch.randn(10, 12, 120, 120)
        crop_size = 64
        
        cropped = dataset.crop_center(img, crop_size, crop_size)
        
        assert cropped.shape == (10, 12, crop_size, crop_size)
    
    def test_crop_center_maintains_channels_and_temporal(self, mock_h5_data):
        """Test that crop_center maintains channel and temporal dimensions."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        channels, temporal = 10, 12
        img = torch.randn(channels, temporal, 120, 120)
        crop_size = 50
        
        cropped = dataset.crop_center(img, crop_size, crop_size)
        
        assert cropped.shape[0] == channels
        assert cropped.shape[1] == temporal
    
    def test_crop_center_asymmetric_crop(self, mock_h5_data):
        """Test center cropping with asymmetric dimensions."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        img = torch.randn(10, 12, 120, 120)
        cropx, cropy = 80, 60
        
        cropped = dataset.crop_center(img, cropx, cropy)
        
        assert cropped.shape == (10, 12, cropy, cropx)


class TestSen4MapDatasetReverseTiling:
    """Test reverse_tiling_pytorch method."""
    
    def test_reverse_tiling_basic(self, mock_h5_data):
        """Test basic reverse tiling."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        # Create a small test tensor
        img = torch.randn(10, 12, 20, 20)
        kernel_size = 3
        
        expanded = dataset.reverse_tiling_pytorch(img, kernel_size)
        
        # Spatial dimensions should be multiplied by kernel_size
        assert expanded.shape == (10, 12, 20 * kernel_size, 20 * kernel_size)
    
    def test_reverse_tiling_maintains_channels(self, mock_h5_data):
        """Test that reverse tiling maintains channel and temporal dimensions."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        channels, temporal = 10, 12
        img = torch.randn(channels, temporal, 20, 20)
        kernel_size = 5
        
        expanded = dataset.reverse_tiling_pytorch(img, kernel_size)
        
        assert expanded.shape[0] == channels
        assert expanded.shape[1] == temporal
        assert expanded.shape[2] == 20 * kernel_size
        assert expanded.shape[3] == 20 * kernel_size
    
    def test_reverse_tiling_kernel_size_5(self, mock_h5_data):
        """Test reverse tiling with kernel size 5."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        img = torch.randn(10, 12, 10, 10)
        kernel_size = 5
        
        expanded = dataset.reverse_tiling_pytorch(img, kernel_size)
        
        assert expanded.shape == (10, 12, 50, 50)
    
    def test_reverse_tiling_assertion_odd_kernel(self, mock_h5_data):
        """Test that even kernel size raises assertion error."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        img = torch.randn(10, 12, 20, 20)
        
        with pytest.raises(AssertionError):
            dataset.reverse_tiling_pytorch(img, kernel_size=4)
    
    def test_reverse_tiling_assertion_minimum_kernel(self, mock_h5_data):
        """Test that kernel size < 3 raises assertion error."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        img = torch.randn(10, 12, 20, 20)
        
        with pytest.raises(AssertionError):
            dataset.reverse_tiling_pytorch(img, kernel_size=1)


class TestSen4MapDatasetMinMaxNormalize:
    """Test min_max_normalize method."""
    
    def test_min_max_normalize_basic(self, mock_h5_data):
        """Test basic min-max normalization."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        # Create test tensor with known values
        tensor = torch.full((10, 12, 20, 20), 500.0)
        q_low = [100.0] * 10
        q_hi = [1000.0] * 10
        
        normalized = dataset.min_max_normalize(tensor, q_low, q_hi)
        
        # Values should be normalized
        assert torch.all(normalized >= 0)
        assert torch.all(normalized <= 1)
    
    def test_min_max_normalize_maintains_shape(self, mock_h5_data):
        """Test that normalization maintains tensor shape."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        original_shape = (10, 12, 20, 20)
        tensor = torch.randn(original_shape)
        q_low = [0.0] * 10
        q_hi = [1.0] * 10
        
        normalized = dataset.min_max_normalize(tensor, q_low, q_hi)
        
        assert normalized.shape == original_shape
    
    def test_min_max_normalize_per_channel(self, mock_h5_data):
        """Test that normalization is applied per channel."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        # Create tensor with different values per channel
        tensor = torch.zeros((3, 12, 20, 20))
        tensor[0, :, :, :] = 500.0
        tensor[1, :, :, :] = 1000.0
        tensor[2, :, :, :] = 1500.0
        
        q_low = [0.0, 500.0, 1000.0]
        q_hi = [1000.0, 1500.0, 2000.0]
        
        normalized = dataset.min_max_normalize(tensor, q_low, q_hi)
        
        # Each channel should be normalized differently
        # Channel 0: (500 - 0) / (1000 - 0) ≈ 0.5
        # Channel 1: (1000 - 500) / (1500 - 500) = 0.5
        # Channel 2: (1500 - 1000) / (2000 - 1000) = 0.5
        assert torch.allclose(normalized[0, 0, 0, 0], torch.tensor(0.5), atol=0.01)
        assert torch.allclose(normalized[1, 0, 0, 0], torch.tensor(0.5), atol=0.01)
        assert torch.allclose(normalized[2, 0, 0, 0], torch.tensor(0.5), atol=0.01)
    
    def test_min_max_normalize_in_place(self, mock_h5_data):
        """Test that normalization modifies tensor in place."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        tensor = torch.full((10, 12, 20, 20), 500.0)
        tensor_id = id(tensor)
        q_low = [100.0] * 10
        q_hi = [1000.0] * 10
        
        normalized = dataset.min_max_normalize(tensor, q_low, q_hi)
        
        # Should modify in place
        assert id(normalized) == tensor_id


class TestSen4MapDatasetLen:
    """Test __len__ method."""
    
    def test_len_default_keys(self, mock_h5_data):
        """Test length with default keys."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        assert len(dataset) == 3
    
    def test_len_custom_keys(self, mock_h5_data):
        """Test length with custom keys."""
        custom_keys = ['patch_0001']
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            h5data_keys=custom_keys
        )
        
        assert len(dataset) == 1
    
    def test_len_empty_keys(self, mock_h5_data):
        """Test length with empty keys list."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            h5data_keys=[]
        )
        
        assert len(dataset) == 0


class TestSen4MapDatasetIntegration:
    """Integration tests combining multiple features."""
    
    def test_full_pipeline_with_all_preprocessing(self, mock_h5_data):
        """Test full data loading pipeline with all preprocessing options."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data,
            crop_size=64,
            resize=True,
            resize_to=[128, 128],
            reverse_tile=True,
            reverse_tile_size=3,
            classification_map="land-cover"
        )
        
        sample = dataset[0]
        
        assert "image" in sample
        assert "label" in sample
        # After crop (64x64), reverse tile (3x), result is 192x192, then resize to 128x128
        assert sample["image"].shape[2] == 128
        assert sample["image"].shape[3] == 128
    
    def test_iteration_over_dataset(self, mock_h5_data):
        """Test iterating over the dataset."""
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5_data
        )
        
        samples = [dataset[i] for i in range(len(dataset))]
        
        assert len(samples) == 3
        for sample in samples:
            assert "image" in sample
            assert "label" in sample
    
    def test_dataset_with_different_labels(self):
        """Test dataset with multiple different label types."""
        mock_h5 = MagicMock(spec=h5py.File)
        
        # Create patches with different labels
        labels = ['A10', 'B11', 'C10', 'F10']
        patches = {}
        
        for idx, label in enumerate(labels):
            n_temporal, height, width = 12, 120, 120
            band_data = {}
            
            for band in ['B2', 'B3', 'B4', 'B5', 'B6', 'B7', 'B8', 'B8A', 'B11', 'B12']:
                band_data[band] = np.random.randint(100, 3000, size=(n_temporal, height, width)).astype(np.float32)
            band_data['SCL'] = np.random.randint(0, 9, size=(n_temporal, height, width))
            
            def make_getitem(data_dict):
                def getitem_patch(key):
                    return data_dict[key]
                return getitem_patch
            
            mock_patch = MagicMock()
            mock_patch.__getitem__.side_effect = make_getitem(band_data)
            
            image_ids = np.array([f'S2_2018{month:02d}15_patch_{idx:03d}' for month in range(1, 13)])
            
            # Create attrs accessor that returns correct values
            def make_attrs_getitem(img_ids, lbl):
                def attrs_getitem(self, key):
                    if key == 'Image_ID':
                        return img_ids
                    elif key == 'lc1':
                        return lbl
                return attrs_getitem
            
            mock_attrs = MagicMock()
            mock_attrs.__getitem__ = make_attrs_getitem(image_ids, label)
            mock_patch.attrs = mock_attrs
            
            patches[f'patch_{idx:04d}'] = mock_patch
        
        mock_h5.keys.return_value = list(patches.keys())
        
        def getitem_h5(key):
            return patches[key]
        
        mock_h5.__getitem__.side_effect = getitem_h5
        
        dataset = Sen4MapDatasetMonthlyComposites(
            h5py_file_object=mock_h5,
            classification_map="land-cover"
        )
        
        # Check that different labels are correctly mapped
        labels_list = [dataset[i]["label"].item() for i in range(len(dataset))]
        
        # A10->0, B11->3, C10->4, F10->1
        expected_labels = [0, 3, 4, 1]
        assert labels_list == expected_labels
