"""
Unit tests for data preprocessing and dataset loading.
"""

import pytest
import torch
import numpy as np
from PIL import Image

import sys
sys.path.insert(0, '.')

from src.data.preprocessing import (
    preprocess_image,
    handle_channels,
    get_optical_transform,
    get_sar_transform
)
from src.data.dataset import (
    CrisisLandMarkDataset,
    create_splits
)


class TestPreprocessing:
    """Test preprocessing functions."""
    
    def test_optical_preprocessing(self):
        """Test optical RGB preprocessing."""
        dummy_image = Image.fromarray(
            np.random.randint(0, 255, (256, 256, 3), dtype=np.uint8)
        )
        
        result = preprocess_image(dummy_image, "optical", size=224)
        
        assert result.shape == (3, 224, 224)
        assert result.dtype == torch.float32
    
    def test_sar_preprocessing(self):
        """Test SAR preprocessing."""
        dummy_image = Image.fromarray(
            np.random.randint(0, 255, (256, 256, 2), dtype=np.uint8)
        )
        
        result = preprocess_image(dummy_image, "sar", size=224)
        
        assert result.shape == (2, 224, 224)
        assert result.dtype == torch.float32
    
    def test_multispectral_preprocessing(self):
        """Test multispectral preprocessing."""
        # Create 12-channel image
        dummy_array = np.random.randint(0, 255, (256, 256, 12), dtype=np.uint8)
        dummy_image = Image.fromarray(dummy_array[..., :3])  # PIL only supports 3 channels
        
        # For 12-channel, we'd need custom handling
        # This test verifies the function accepts the modality
        result = preprocess_image(dummy_image, "optical", size=224)
        
        assert result.shape[0] == 3  # Should be 3 channels from RGB
    
    def test_channel_handling(self):
        """Test channel mismatch handling."""
        # Create 4-channel image (should be trimmed to 3 for optical)
        image_4ch = np.random.randint(0, 255, (64, 64, 4), dtype=np.uint8)
        
        result = handle_channels(image_4ch, 3, "optical")
        
        assert result.shape[-1] == 3
    
    def test_invalid_modality(self):
        """Test invalid modality raises error."""
        dummy_image = Image.fromarray(
            np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)
        )
        
        with pytest.raises(ValueError):
            preprocess_image(dummy_image, "invalid_modality")


class TestDataset:
    """Test dataset class."""
    
    def test_dataset_creation(self):
        """Test dataset can be created."""
        dataset = CrisisLandMarkDataset(modality="optical")
        
        assert len(dataset) > 0
    
    def test_dataset_getitem(self):
        """Test dataset returns correct format."""
        dataset = CrisisLandMarkDataset(modality="optical")
        
        image, modality_label, class_label = dataset[0]
        
        assert isinstance(image, torch.Tensor)
        assert isinstance(modality_label, int)
        assert isinstance(class_label, int)
    
    def test_modality_labels(self):
        """Test modality labels are correct."""
        # Test optical and SAR (multispectral requires 12-channel images)
        for modality, expected_label in [("optical", 0), ("sar", 1)]:
            dataset = CrisisLandMarkDataset(modality=modality)
            _, modality_label, _ = dataset[0]
            
            assert modality_label == expected_label


class TestSplitting:
    """Test data splitting."""
    
    def test_split_no_overlap(self):
        """Test query/gallery split has no overlap."""
        dataset = CrisisLandMarkDataset(modality="optical")
        
        query_idx, gallery_idx = create_splits(dataset, query_ratio=0.2)
        
        # Check no overlap
        assert len(set(query_idx) & set(gallery_idx)) == 0
    
    def test_split_ratio(self):
        """Test split maintains approximate ratio."""
        dataset = CrisisLandMarkDataset(modality="optical")
        
        query_idx, gallery_idx = create_splits(dataset, query_ratio=0.2)
        
        total = len(query_idx) + len(gallery_idx)
        query_ratio = len(query_idx) / total
        
        # Allow some tolerance
        assert 0.15 <= query_ratio <= 0.25
    
    def test_split_reproducibility(self):
        """Test split is reproducible with same seed."""
        dataset = CrisisLandMarkDataset(modality="optical")
        
        query1, gallery1 = create_splits(dataset, seed=42)
        query2, gallery2 = create_splits(dataset, seed=42)
        
        assert query1 == query2
        assert gallery1 == gallery2


if __name__ == "__main__":
    pytest.main([__file__, "-v"])