tests/test_dataset.py

"""
Unit Tests for VintageGAN Dataset Module
Specification Reference: Day 1-2 - Write unit tests for data loading

Tests cover:
1. ImageNetSubsetDataset functionality
2. VintageGANDataset with dummy defect generator
3. DataLoader creation and validation
4. Tensor shape and value range verification
5. Augmentation pipeline

Author: VintageGAN Project
Date: 2024
"""

import os
import sys
from pathlib import Path
import tempfile
import shutil

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

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))

from training.dataset import (
    ImageNetSubsetDataset,
    VintageGANDataset,
    FilmSetDataset,
    create_dataloaders,
    validate_dataloader,
)
from training.download_data import create_dummy_dataset


class TestImageNetSubsetDataset:
    """Test suite for ImageNetSubsetDataset class."""

    @pytest.fixture
    def dummy_dataset_dir(self, tmp_path):
        """Create temporary dummy dataset for testing."""
        dataset_dir = tmp_path / "imagenet_subset"
        create_dummy_dataset(str(dataset_dir), num_train=10, num_val=5, image_size=512)
        return dataset_dir

    def test_dataset_initialization(self, dummy_dataset_dir):
        """Test dataset can be initialized correctly."""
        dataset = ImageNetSubsetDataset(
            root_dir=str(dummy_dataset_dir), split="train", image_size=512
        )

        assert len(dataset) > 0, "Dataset should contain images"
        assert dataset.image_size == 512, "Image size should be 512"
        assert dataset.split == "train", "Split should be 'train'"

    def test_dataset_length(self, dummy_dataset_dir):
        """Test dataset returns correct length."""
        train_dataset = ImageNetSubsetDataset(
            root_dir=str(dummy_dataset_dir), split="train"
        )
        val_dataset = ImageNetSubsetDataset(
            root_dir=str(dummy_dataset_dir), split="val"
        )

        assert (
            len(train_dataset) == 10
        ), f"Expected 10 train images, got {len(train_dataset)}"
        assert len(val_dataset) == 5, f"Expected 5 val images, got {len(val_dataset)}"

    def test_getitem_returns_correct_format(self, dummy_dataset_dir):
        """Test __getitem__ returns (tensor, path) tuple."""
        dataset = ImageNetSubsetDataset(root_dir=str(dummy_dataset_dir), split="train")

        image, path = dataset[0]

        assert isinstance(image, torch.Tensor), "Image should be torch.Tensor"
        assert isinstance(path, str), "Path should be string"
        assert image.dim() == 3, f"Image should be 3D tensor, got {image.dim()}D"

    def test_image_shape(self, dummy_dataset_dir):
        """Test images have correct shape (3, 512, 512)."""
        dataset = ImageNetSubsetDataset(
            root_dir=str(dummy_dataset_dir), split="train", image_size=512
        )

        image, _ = dataset[0]

        assert image.shape == (
            3,
            512,
            512,
        ), f"Expected shape (3, 512, 512), got {image.shape}"

    def test_image_value_range(self, dummy_dataset_dir):
        """Test images are normalized to [-1, 1]."""
        dataset = ImageNetSubsetDataset(root_dir=str(dummy_dataset_dir), split="train")

        image, _ = dataset[0]

        assert (
            image.min() >= -1.5
        ), f"Image min value {image.min()} below expected range"
        assert image.max() <= 1.5, f"Image max value {image.max()} above expected range"
        # Most values should be in [-1, 1], allowing some tolerance for numerical precision

    def test_invalid_split_raises_error(self, dummy_dataset_dir):
        """Test invalid split parameter raises ValueError."""
        with pytest.raises(ValueError, match="Split must be 'train' or 'val'"):
            ImageNetSubsetDataset(
                root_dir=str(dummy_dataset_dir), split="test"  # Invalid split
            )

    def test_get_sample_images(self, dummy_dataset_dir):
        """Test get_sample_images returns correct number of samples."""
        dataset = ImageNetSubsetDataset(root_dir=str(dummy_dataset_dir), split="train")

        samples = dataset.get_sample_images(num_samples=5)

        assert len(samples) == 5, f"Expected 5 samples, got {len(samples)}"
        assert all(
            isinstance(img, torch.Tensor) for img in samples
        ), "All samples should be tensors"


class TestVintageGANDataset:
    """Test suite for VintageGANDataset class."""

    @pytest.fixture
    def dummy_clean_dataset(self, tmp_path):
        """Create dummy clean dataset."""
        dataset_dir = tmp_path / "imagenet_subset"
        create_dummy_dataset(str(dataset_dir), num_train=5, num_val=2)
        return ImageNetSubsetDataset(
            root_dir=str(dataset_dir), split="train", image_size=512
        )

    @pytest.fixture
    def dummy_defect_generator(self):
        """Create dummy defect generator function."""

        def generate_defects(image: np.ndarray, conditions: np.ndarray) -> np.ndarray:
            """
            Dummy defect generator that adds small noise.

            Args:
                image: (H, W, 3) uint8 array in [0, 255]
                conditions: (6,) float array in [0, 1]

            Returns:
                Defected image (H, W, 3) uint8 in [0, 255]
            """
            # Add small noise proportional to condition intensity
            noise_intensity = np.mean(conditions) * 10
            noise = np.random.randn(*image.shape) * noise_intensity
            defected = image.astype(np.float32) + noise
            defected = np.clip(defected, 0, 255).astype(np.uint8)
            return defected

        return generate_defects

    def test_dataset_initialization(self, dummy_clean_dataset, dummy_defect_generator):
        """Test VintageGANDataset initialization."""
        dataset = VintageGANDataset(
            clean_dataset=dummy_clean_dataset,
            defect_generator=dummy_defect_generator,
            num_variants=3,
            return_clean=True,
        )

        expected_size = len(dummy_clean_dataset) * 3
        assert (
            len(dataset) == expected_size
        ), f"Expected {expected_size} pairs, got {len(dataset)}"

    def test_getitem_returns_dict(self, dummy_clean_dataset, dummy_defect_generator):
        """Test __getitem__ returns dictionary with correct keys."""
        dataset = VintageGANDataset(
            clean_dataset=dummy_clean_dataset,
            defect_generator=dummy_defect_generator,
            num_variants=2,
            return_clean=True,
        )

        batch = dataset[0]

        assert isinstance(batch, dict), "Output should be dictionary"
        assert "clean" in batch, "Dictionary should contain 'clean' key"
        assert "defected" in batch, "Dictionary should contain 'defected' key"
        assert "condition" in batch, "Dictionary should contain 'condition' key"
        assert "image_id" in batch, "Dictionary should contain 'image_id' key"

    def test_output_shapes(self, dummy_clean_dataset, dummy_defect_generator):
        """Test output tensors have correct shapes."""
        dataset = VintageGANDataset(
            clean_dataset=dummy_clean_dataset,
            defect_generator=dummy_defect_generator,
            num_variants=2,
            return_clean=True,
        )

        batch = dataset[0]

        assert batch["clean"].shape == (
            3,
            512,
            512,
        ), f"Clean shape incorrect: {batch['clean'].shape}"
        assert batch["defected"].shape == (
            3,
            512,
            512,
        ), f"Defected shape incorrect: {batch['defected'].shape}"
        assert batch["condition"].shape == (
            6,
        ), f"Condition shape incorrect: {batch['condition'].shape}"

    def test_condition_value_range(self, dummy_clean_dataset, dummy_defect_generator):
        """Test condition vector is in [0, 1] range."""
        dataset = VintageGANDataset(
            clean_dataset=dummy_clean_dataset,
            defect_generator=dummy_defect_generator,
            num_variants=2,
        )

        batch = dataset[0]
        condition = batch["condition"]

        assert torch.all(condition >= 0.0), f"Condition min {condition.min()} below 0"
        assert torch.all(condition <= 1.0), f"Condition max {condition.max()} above 1"

    def test_reproducibility_with_seed(
        self, dummy_clean_dataset, dummy_defect_generator
    ):
        """Test same index produces same conditions (reproducibility)."""
        dataset = VintageGANDataset(
            clean_dataset=dummy_clean_dataset,
            defect_generator=dummy_defect_generator,
            num_variants=3,
        )

        # Get same item twice
        batch1 = dataset[2]
        batch2 = dataset[2]

        # Conditions should be identical (deterministic based on image+variant)
        assert torch.allclose(
            batch1["condition"], batch2["condition"], atol=1e-6
        ), "Same index should produce same conditions"

    def test_different_variants_have_different_conditions(
        self, dummy_clean_dataset, dummy_defect_generator
    ):
        """Test different variants of same image have different conditions."""
        dataset = VintageGANDataset(
            clean_dataset=dummy_clean_dataset,
            defect_generator=dummy_defect_generator,
            num_variants=5,
        )

        # Get different variants of first clean image (indices 0, 1, 2, 3, 4)
        conditions = [dataset[i]["condition"] for i in range(5)]

        # Check they're different from each other
        for i in range(len(conditions)):
            for j in range(i + 1, len(conditions)):
                assert not torch.allclose(
                    conditions[i], conditions[j], atol=1e-3
                ), f"Variants {i} and {j} have identical conditions"

    def test_return_clean_false(self, dummy_clean_dataset, dummy_defect_generator):
        """Test return_clean=False doesn't include clean image."""
        dataset = VintageGANDataset(
            clean_dataset=dummy_clean_dataset,
            defect_generator=dummy_defect_generator,
            num_variants=2,
            return_clean=False,
        )

        batch = dataset[0]

        assert (
            "clean" not in batch
        ), "Should not include 'clean' when return_clean=False"
        assert "defected" in batch, "Should still include 'defected'"
        assert "condition" in batch, "Should still include 'condition'"


class TestDataLoaderCreation:
    """Test suite for dataloader creation and validation."""

    @pytest.fixture
    def config_file(self, tmp_path):
        """Create temporary config file."""
        # Create dummy dataset
        dataset_dir = tmp_path / "data" / "imagenet_subset"
        create_dummy_dataset(str(dataset_dir), num_train=8, num_val=4)

        # Create config file
        config_path = tmp_path / "config.yaml"
        config_content = f"""
dataset:
  imagenet_path: "{dataset_dir}"
  image_size: 512
  variants_per_image: 2
  val_variants_per_image: 1

hardware:
  batch_size: 2
  num_workers: 0
  pin_memory: false

augmentation:
  horizontal_flip_prob: 0.5
  brightness_jitter: 0.1
  contrast_jitter: 0.1
"""
        config_path.write_text(config_content)
        return config_path

    def test_create_dataloaders_without_defect_generator(self, config_file):
        """Test creating dataloaders without defect generator (clean images only)."""
        dataloaders = create_dataloaders(str(config_file), defect_generator=None)

        assert "train" in dataloaders, "Should contain 'train' dataloader"
        assert "val" in dataloaders, "Should contain 'val' dataloader"

        # Test loading a batch
        train_batch = next(iter(dataloaders["train"]))
        assert len(train_batch) == 2, "Should return tuple of (images, paths)"

    def test_create_dataloaders_with_defect_generator(self, config_file):
        """Test creating dataloaders with defect generator."""

        def dummy_defect_gen(image, conditions):
            """Dummy defect generator."""
            return image.copy()

        dataloaders = create_dataloaders(
            str(config_file), defect_generator=dummy_defect_gen
        )

        # Test loading a batch
        train_batch = next(iter(dataloaders["train"]))

        assert isinstance(
            train_batch, dict
        ), "Should return dictionary with VintageGAN dataset"
        assert "clean" in train_batch, "Should contain clean images"
        assert "defected" in train_batch, "Should contain defected images"
        assert "condition" in train_batch, "Should contain conditions"

    def test_batch_shapes(self, config_file):
        """Test batch tensors have correct shapes."""

        def dummy_defect_gen(image, conditions):
            return image.copy()

        dataloaders = create_dataloaders(
            str(config_file), defect_generator=dummy_defect_gen
        )
        train_batch = next(iter(dataloaders["train"]))

        batch_size = 2
        assert train_batch["clean"].shape == (
            batch_size,
            3,
            512,
            512,
        ), f"Clean batch shape incorrect: {train_batch['clean'].shape}"
        assert train_batch["defected"].shape == (
            batch_size,
            3,
            512,
            512,
        ), f"Defected batch shape incorrect: {train_batch['defected'].shape}"
        assert train_batch["condition"].shape == (
            batch_size,
            6,
        ), f"Condition batch shape incorrect: {train_batch['condition'].shape}"


class TestTensorConversion:
    """Test tensor conversion utilities."""

    @pytest.fixture
    def sample_dataset(self, tmp_path):
        """Create sample dataset for testing."""
        dataset_dir = tmp_path / "imagenet_subset"
        create_dummy_dataset(str(dataset_dir), num_train=2, num_val=1)

        clean_dataset = ImageNetSubsetDataset(root_dir=str(dataset_dir), split="train")

        def dummy_defect_gen(image, conditions):
            return image.copy()

        return VintageGANDataset(
            clean_dataset=clean_dataset,
            defect_generator=dummy_defect_gen,
            num_variants=1,
        )

    def test_tensor_to_numpy_conversion(self, sample_dataset):
        """Test _tensor_to_numpy preserves value ranges."""
        # Create a known tensor
        tensor = torch.zeros(3, 512, 512)  # All zeros
        tensor[0, :, :] = -1.0  # R channel = -1 (should become 0)
        tensor[1, :, :] = 0.0  # G channel = 0 (should become 127-128)
        tensor[2, :, :] = 1.0  # B channel = 1 (should become 255)

        # Convert using dataset method
        array = sample_dataset._tensor_to_numpy(tensor)

        assert array.shape == (512, 512, 3), f"Shape mismatch: {array.shape}"
        assert array.dtype == np.uint8, f"Dtype should be uint8, got {array.dtype}"
        assert array[0, 0, 0] == 0, f"R channel should be 0, got {array[0, 0, 0]}"
        assert (
            120 <= array[0, 0, 1] <= 135
        ), f"G channel should be ~128, got {array[0, 0, 1]}"
        assert array[0, 0, 2] == 255, f"B channel should be 255, got {array[0, 0, 2]}"

    def test_numpy_to_tensor_conversion(self, sample_dataset):
        """Test _numpy_to_tensor preserves value ranges."""
        # Create a known array
        array = np.zeros((512, 512, 3), dtype=np.uint8)
        array[:, :, 0] = 0  # R = 0 (should become -1)
        array[:, :, 1] = 128  # G = 128 (should become ~0)
        array[:, :, 2] = 255  # B = 255 (should become 1)

        # Convert using dataset method
        tensor = sample_dataset._numpy_to_tensor(array)

        assert tensor.shape == (3, 512, 512), f"Shape mismatch: {tensor.shape}"
        assert (
            torch.abs(tensor[0, 0, 0] - (-1.0)) < 0.02
        ), f"R channel should be -1, got {tensor[0, 0, 0]}"
        assert (
            torch.abs(tensor[1, 0, 0] - 0.0) < 0.02
        ), f"G channel should be 0, got {tensor[1, 0, 0]}"
        assert (
            torch.abs(tensor[2, 0, 0] - 1.0) < 0.02
        ), f"B channel should be 1, got {tensor[2, 0, 0]}"

    def test_roundtrip_conversion(self, sample_dataset):
        """Test tensor -> numpy -> tensor roundtrip."""
        original = torch.rand(3, 512, 512) * 1.8 - 0.9

        # Roundtrip conversion
        array = sample_dataset._tensor_to_numpy(original)
        recovered = sample_dataset._numpy_to_tensor(array)

        # Should be close (some loss due to uint8 quantization)
        assert recovered.shape == original.shape, "Shape should be preserved"
        # Allow some error due to quantization
        max_error = torch.abs(recovered - original).max()
        assert max_error < 0.02, f"Roundtrip error too large: {max_error}"


def test_validate_dataloader(tmp_path, capsys):
    """Test validate_dataloader function."""
    # Create dummy dataset
    dataset_dir = tmp_path / "imagenet_subset"
    create_dummy_dataset(str(dataset_dir), num_train=8, num_val=4)

    # Create config
    config_path = tmp_path / "config.yaml"
    config_content = f"""
dataset:
  imagenet_path: "{dataset_dir}"
  image_size: 512
  variants_per_image: 1
  val_variants_per_image: 1

hardware:
  batch_size: 2
  num_workers: 0
  pin_memory: false
"""
    config_path.write_text(config_content)

    # Create dataloader
    def dummy_defect_gen(image, conditions):
        return image.copy()

    dataloaders = create_dataloaders(
        str(config_path), defect_generator=dummy_defect_gen
    )

    # Validate (should not raise)
    validate_dataloader(dataloaders["train"], num_batches=2)

    # Check output was printed
    captured = capsys.readouterr()
    assert "DATALOADER VALIDATION" in captured.out
    assert "All validation checks passed" in captured.out


if __name__ == "__main__":
    """Run tests with pytest."""
    pytest.main([__file__, "-v", "--tb=short"])