"""Comprehensive tests for Landslide4SenseNonGeo dataset."""

import pytest
import numpy as np
import torch
import h5py
import albumentations as A
from pathlib import Path
from unittest.mock import Mock, patch, MagicMock
from matplotlib.figure import Figure

from terratorch.datasets.landslide4sense import Landslide4SenseNonGeo


@pytest.fixture
def mock_data_root(tmp_path):
    """Create a mock data directory structure."""
    # Create directory structure
    for split in ["train", "validation", "test"]:
        (tmp_path / "images" / split).mkdir(parents=True, exist_ok=True)
        (tmp_path / "annotations" / split).mkdir(parents=True, exist_ok=True)
    
    # Create mock HDF5 files
    for split in ["train", "validation", "test"]:
        for i in range(3):
            # Create image file
            image_path = tmp_path / "images" / split / f"image_{i:04d}.h5"
            with h5py.File(image_path, "w") as f:
                # 14 bands: 12 spectral + slope + DEM
                image_data = np.random.rand(128, 128, 14).astype(np.float32)
                f.create_dataset("img", data=image_data)
            
            # Create mask file
            mask_path = tmp_path / "annotations" / split / f"mask_{i:04d}.h5"
            with h5py.File(mask_path, "w") as f:
                # Binary mask (landslide vs no landslide)
                mask_data = np.random.randint(0, 2, size=(128, 128), dtype=np.uint8)
                f.create_dataset("mask", data=mask_data)
    
    return tmp_path


class TestLandslide4SenseInitialization:
    """Test dataset initialization and parameter validation."""
    
    def test_basic_initialization(self, mock_data_root):
        """Test basic dataset initialization."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        
        assert dataset.split == "train"
        assert len(dataset.bands) == 14  # Default: all bands
        assert len(dataset) == 3
        assert len(dataset.band_indices) == 14
    
    def test_initialization_with_validation_split(self, mock_data_root):
        """Test initialization with validation split."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="val")
        
        assert dataset.split == "val"
        assert len(dataset) == 3
    
    def test_initialization_with_test_split(self, mock_data_root):
        """Test initialization with test split."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="test")
        
        assert dataset.split == "test"
        assert len(dataset) == 3
    
    def test_initialization_with_invalid_split(self, mock_data_root):
        """Test that invalid split raises ValueError."""
        with pytest.raises(ValueError, match="Incorrect split"):
            Landslide4SenseNonGeo(data_root=str(mock_data_root), split="invalid")
    
    def test_initialization_with_rgb_bands(self, mock_data_root):
        """Test initialization with RGB bands only."""
        dataset = Landslide4SenseNonGeo(
            data_root=str(mock_data_root),
            split="train",
            bands=Landslide4SenseNonGeo.BAND_SETS["rgb"]
        )
        
        assert len(dataset.bands) == 3
        assert dataset.bands == ("RED", "GREEN", "BLUE")
        assert len(dataset.band_indices) == 3
        assert dataset.band_indices == [3, 2, 1]  # RED, GREEN, BLUE indices
    
    def test_initialization_with_custom_bands(self, mock_data_root):
        """Test initialization with custom band selection."""
        custom_bands = ("RED", "NIR_BROAD", "SWIR_1")
        dataset = Landslide4SenseNonGeo(
            data_root=str(mock_data_root),
            split="train",
            bands=custom_bands
        )
        
        assert dataset.bands == custom_bands
        assert len(dataset.band_indices) == 3
        assert dataset.band_indices == [3, 7, 10]  # RED, NIR_BROAD, SWIR_1
    
    def test_initialization_with_invalid_bands(self, mock_data_root):
        """Test that invalid bands raise ValueError."""
        with pytest.raises(ValueError):
            Landslide4SenseNonGeo(
                data_root=str(mock_data_root),
                split="train",
                bands=("INVALID_BAND",)
            )
    
    def test_initialization_with_custom_transform(self, mock_data_root):
        """Test initialization with custom transform."""
        transform = A.Compose([A.HorizontalFlip(p=1.0)])
        dataset = Landslide4SenseNonGeo(
            data_root=str(mock_data_root),
            split="train",
            transform=transform
        )
        
        assert dataset.transform == transform
    
    def test_band_sets_attribute(self):
        """Test BAND_SETS class attribute."""
        assert "all" in Landslide4SenseNonGeo.BAND_SETS
        assert "rgb" in Landslide4SenseNonGeo.BAND_SETS
        assert len(Landslide4SenseNonGeo.BAND_SETS["all"]) == 14
        assert len(Landslide4SenseNonGeo.BAND_SETS["rgb"]) == 3
    
    def test_all_band_names_attribute(self):
        """Test all_band_names contains expected bands."""
        expected_bands = [
            "COASTAL AEROSOL", "BLUE", "GREEN", "RED", "RED_EDGE_1",
            "RED_EDGE_2", "RED_EDGE_3", "NIR_BROAD", "WATER_VAPOR",
            "CIRRUS", "SWIR_1", "SWIR_2", "SLOPE", "DEM"
        ]
        assert Landslide4SenseNonGeo.all_band_names == tuple(expected_bands)
    
    def test_rgb_bands_attribute(self):
        """Test rgb_bands contains RGB channels."""
        assert Landslide4SenseNonGeo.rgb_bands == ("RED", "GREEN", "BLUE")


class TestLandslide4SenseGetItem:
    """Test dataset __getitem__ functionality."""
    
    def test_getitem_returns_dict(self, mock_data_root):
        """Test that __getitem__ returns a dictionary."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        sample = dataset[0]
        
        assert isinstance(sample, dict)
        assert "image" in sample
        assert "mask" in sample
    
    def test_getitem_image_shape(self, mock_data_root):
        """Test that image has correct shape with all bands."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        sample = dataset[0]
        
        assert sample["image"].shape[0] == 14  # All bands
        assert sample["image"].shape[1] == 128  # Height
        assert sample["image"].shape[2] == 128  # Width
    
    def test_getitem_image_shape_with_rgb(self, mock_data_root):
        """Test that image has correct shape with RGB bands only."""
        dataset = Landslide4SenseNonGeo(
            data_root=str(mock_data_root),
            split="train",
            bands=Landslide4SenseNonGeo.BAND_SETS["rgb"]
        )
        sample = dataset[0]
        
        assert sample["image"].shape[0] == 3  # RGB only
        assert sample["image"].shape[1] == 128
        assert sample["image"].shape[2] == 128
    
    def test_getitem_mask_shape(self, mock_data_root):
        """Test that mask has correct shape."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        sample = dataset[0]
        
        assert sample["mask"].shape == torch.Size([128, 128])
    
    def test_getitem_mask_dtype(self, mock_data_root):
        """Test that mask is converted to long tensor."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        sample = dataset[0]
        
        assert sample["mask"].dtype == torch.long
    
    def test_getitem_image_dtype(self, mock_data_root):
        """Test that image is float32 tensor."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        sample = dataset[0]
        
        assert sample["image"].dtype == torch.float32
    
    def test_getitem_with_transform(self, mock_data_root):
        """Test that transforms are applied."""
        transform = A.Compose([A.HorizontalFlip(p=1.0)])
        dataset = Landslide4SenseNonGeo(
            data_root=str(mock_data_root),
            split="train",
            transform=transform
        )
        sample = dataset[0]
        
        # Should have flipped the image
        assert "image" in sample
        assert "mask" in sample
    
    def test_getitem_different_indices(self, mock_data_root):
        """Test accessing different samples."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        
        sample_0 = dataset[0]
        sample_1 = dataset[1]
        sample_2 = dataset[2]
        
        assert sample_0["image"].shape == sample_1["image"].shape == sample_2["image"].shape
        assert sample_0["mask"].shape == sample_1["mask"].shape == sample_2["mask"].shape
    
    def test_getitem_band_selection(self, mock_data_root):
        """Test that band selection is correctly applied."""
        bands = ("RED", "GREEN", "BLUE", "NIR_BROAD")
        dataset = Landslide4SenseNonGeo(
            data_root=str(mock_data_root),
            split="train",
            bands=bands
        )
        sample = dataset[0]
        
        assert sample["image"].shape[0] == 4  # 4 selected bands


class TestLandslide4SenseLen:
    """Test dataset length functionality."""
    
    def test_len_returns_correct_count(self, mock_data_root):
        """Test that __len__ returns correct number of samples."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        assert len(dataset) == 3
    
    def test_len_for_validation_split(self, mock_data_root):
        """Test length for validation split."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="val")
        assert len(dataset) == 3
    
    def test_len_for_test_split(self, mock_data_root):
        """Test length for test split."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="test")
        assert len(dataset) == 3


class TestLandslide4SensePlot:
    """Test dataset plotting functionality."""
    
    def test_plot_returns_figure(self, mock_data_root):
        """Test that plot returns a matplotlib Figure."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        sample = dataset[0]
        
        fig = dataset.plot(sample)
        assert isinstance(fig, Figure)
    
    def test_plot_with_suptitle(self, mock_data_root):
        """Test plot with custom suptitle."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        sample = dataset[0]
        
        fig = dataset.plot(sample, suptitle="Test Sample")
        assert isinstance(fig, Figure)
    
    def test_plot_with_prediction(self, mock_data_root):
        """Test plot with prediction mask."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        sample = dataset[0]
        
        # Add prediction to sample
        sample["prediction"] = torch.randint(0, 2, (128, 128))
        
        fig = dataset.plot(sample)
        assert isinstance(fig, Figure)
        # Should create 4 subplots when prediction is present
        assert len(fig.axes) == 4
    
    def test_plot_without_prediction(self, mock_data_root):
        """Test plot without prediction mask."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        sample = dataset[0]
        
        fig = dataset.plot(sample)
        assert isinstance(fig, Figure)
        # Should create 3 subplots when no prediction
        assert len(fig.axes) == 3
    
    def test_plot_with_class_names(self, mock_data_root):
        """Test plot with class names."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        sample = dataset[0]
        sample["class_names"] = ["No Landslide", "Landslide"]
        
        fig = dataset.plot(sample)
        assert isinstance(fig, Figure)
    
    def test_plot_with_rgb_bands_only(self, mock_data_root):
        """Test plot with RGB bands."""
        dataset = Landslide4SenseNonGeo(
            data_root=str(mock_data_root),
            split="train",
            bands=Landslide4SenseNonGeo.BAND_SETS["rgb"]
        )
        sample = dataset[0]
        
        fig = dataset.plot(sample)
        assert isinstance(fig, Figure)
    
    def test_plot_missing_rgb_bands_raises_error(self, mock_data_root):
        """Test that plot raises error when RGB bands are missing."""
        dataset = Landslide4SenseNonGeo(
            data_root=str(mock_data_root),
            split="train",
            bands=("SWIR_1", "SWIR_2", "NIR_BROAD")  # No RGB bands
        )
        sample = dataset[0]
        
        with pytest.raises(ValueError, match="doesn't contain some of the RGB bands"):
            dataset.plot(sample)
    
    def test_plot_with_tensor_image(self, mock_data_root):
        """Test plot handles tensor images correctly."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        sample = dataset[0]
        
        # Ensure image is a tensor (should be by default)
        assert torch.is_tensor(sample["image"])
        
        fig = dataset.plot(sample)
        assert isinstance(fig, Figure)
    
    def test_plot_with_numpy_image(self, mock_data_root):
        """Test plot handles numpy images correctly."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        sample = dataset[0]
        
        # Convert image to numpy
        sample["image"] = sample["image"].numpy()
        
        fig = dataset.plot(sample)
        assert isinstance(fig, Figure)


class TestLandslide4SenseIntegration:
    """Integration tests for the dataset."""
    
    def test_full_pipeline(self, mock_data_root):
        """Test full data loading pipeline."""
        dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        
        # Test iteration over dataset
        samples = [dataset[i] for i in range(len(dataset))]
        
        assert len(samples) == 3
        for sample in samples:
            assert "image" in sample
            assert "mask" in sample
            assert sample["image"].shape[0] == 14
            assert sample["mask"].dtype == torch.long
    
    def test_different_splits_have_separate_data(self, mock_data_root):
        """Test that different splits load different data."""
        train_dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="train")
        val_dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="val")
        test_dataset = Landslide4SenseNonGeo(data_root=str(mock_data_root), split="test")
        
        assert len(train_dataset) == len(val_dataset) == len(test_dataset)
        
        # Each should access different files
        train_files = train_dataset.image_files
        val_files = val_dataset.image_files
        test_files = test_dataset.image_files
        
        assert str(train_files[0].parent) != str(val_files[0].parent)
        assert str(val_files[0].parent) != str(test_files[0].parent)
    
    def test_iteration_with_augmentation(self, mock_data_root):
        """Test dataset iteration with augmentations."""
        transform = A.Compose([
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.5),
            A.RandomRotate90(p=0.5)
        ])
        dataset = Landslide4SenseNonGeo(
            data_root=str(mock_data_root),
            split="train",
            transform=transform
        )
        
        for i in range(len(dataset)):
            sample = dataset[i]
            assert sample["image"].shape == torch.Size([14, 128, 128])
            assert sample["mask"].shape == torch.Size([128, 128])
    
    def test_band_consistency_across_samples(self, mock_data_root):
        """Test that band selection is consistent across all samples."""
        bands = ("RED", "GREEN", "BLUE")
        dataset = Landslide4SenseNonGeo(
            data_root=str(mock_data_root),
            split="train",
            bands=bands
        )
        
        for i in range(len(dataset)):
            sample = dataset[i]
            assert sample["image"].shape[0] == 3  # RGB bands
