src/data/dataset.py

"""
Dataset loading and splitting for cross-modal retrieval.

Handles:
- Loading CrisisLandMark dataset
- Per-modality preprocessing
- Query/gallery splitting (80/20)
- Ground-truth label preparation
"""

import torch
from torch.utils.data import Dataset, DataLoader
from pathlib import Path
from typing import Tuple, List, Dict, Optional
import numpy as np
from PIL import Image
from sklearn.model_selection import train_test_split

from .preprocessing import preprocess_image, handle_channels


class CrisisLandMarkDataset(Dataset):
    """
    Dataset class for CrisisLandMark satellite imagery.
    
    Supports:
    - Optical (Sentinel-2 RGB)
    - SAR (Sentinel-1 VV/VH)
    - Multispectral (Sentinel-2 all bands)
    """
    
    def __init__(
        self,
        data_dir: str = "data/raw/crisislandmark",
        modality: str = "optical",
        split: str = "train",
        transform=None,
        size: int = 224
    ):
        """
        Initialize dataset.
        
        Args:
            data_dir: Path to dataset
            modality: "optical", "sar", or "multispectral"
            split: "train", "validation", or "test"
            transform: Optional custom transform
            size: Image resize size
        """
        self.data_dir = Path(data_dir)
        self.modality = modality
        self.split = split
        self.size = size
        self.transform = transform
        
        # ponytail: placeholder - load from actual dataset
        # Real implementation would load from HuggingFace datasets
        self.samples = self._load_samples()
        self.labels = self._load_labels()
    
    def _load_samples(self) -> List[Dict]:
        """Load sample metadata."""
        # Placeholder - will be replaced with actual data loading
        return [{"id": i, "path": f"sample_{i}.png"} for i in range(100)]
    
    def _load_labels(self) -> Dict[int, int]:
        """Load ground-truth labels."""
        # Placeholder - will be replaced with actual labels
        return {i: i % 10 for i in range(100)}
    
    def __len__(self) -> int:
        return len(self.samples)
    
    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int, int]:
        """
        Get sample.
        
        Returns:
            (image_tensor, modality_label, class_label)
        """
        sample = self.samples[idx]
        
        # Load image (placeholder)
        image = Image.fromarray(np.random.randint(0, 255, (256, 256, 3), dtype=np.uint8))
        
        # Preprocess
        if self.transform:
            image_tensor = self.transform(image)
        else:
            image_tensor = preprocess_image(image, self.modality, self.size)
        
        # Modality label (0=optical, 1=sar, 2=multispectral)
        modality_label = {"optical": 0, "sar": 1, "multispectral": 2}[self.modality]
        
        # Class label
        class_label = self.labels.get(idx, 0)
        
        return image_tensor, modality_label, class_label


def create_splits(
    dataset: CrisisLandMarkDataset,
    query_ratio: float = 0.2,
    seed: int = 42
) -> Tuple[List[int], List[int]]:
    """
    Create query/gallery split with no overlap.
    
    Args:
        dataset: Full dataset
        query_ratio: Fraction for query set
        seed: Random seed for reproducibility
        
    Returns:
        (query_indices, gallery_indices)
    """
    indices = list(range(len(dataset)))
    
    # Stratify by class label if available
    labels = [dataset.labels.get(i, 0) for i in indices]
    
    query_idx, gallery_idx = train_test_split(
        indices,
        test_size=1 - query_ratio,
        random_state=seed,
        stratify=labels
    )
    
    # Verify no overlap
    assert len(set(query_idx) & set(gallery_idx)) == 0, "Query and gallery sets overlap!"
    
    return query_idx, gallery_idx


def get_dataloaders(
    data_dir: str = "data/raw/crisislandmark",
    modality: str = "optical",
    batch_size: int = 32,
    size: int = 224,
    num_workers: int = 4
) -> Tuple[DataLoader, DataLoader]:
    """
    Get train/test dataloaders.
    
    Args:
        data_dir: Path to dataset
        modality: Modality type
        batch_size: Batch size
        size: Image size
        num_workers: Number of workers
        
    Returns:
        (train_loader, test_loader)
    """
    train_dataset = CrisisLandMarkDataset(data_dir, modality, "train", size=size)
    test_dataset = CrisisLandMarkDataset(data_dir, modality, "test", size=size)
    
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers
    )
    
    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers
    )
    
    return train_loader, test_loader


# Self-check
if __name__ == "__main__":
    # Test dataset creation
    dataset = CrisisLandMarkDataset(modality="optical")
    
    # Test split
    query_idx, gallery_idx = create_splits(dataset, query_ratio=0.2)
    
    print(f"Total samples: {len(dataset)}")
    print(f"Query set: {len(query_idx)} samples")
    print(f"Gallery set: {len(gallery_idx)} samples")
    print(f"Overlap: {len(set(query_idx) & set(gallery_idx))} (should be 0)")
    
    # Test dataloader
    train_loader, test_loader = get_dataloaders(modality="optical", batch_size=4)
    batch = next(iter(train_loader))
    
    print(f"\nBatch shapes:")
    print(f"  Images: {batch[0].shape}")
    print(f"  Modality labels: {batch[1]}")
    print(f"  Class labels: {batch[2]}")
    
    print("\nDataset test passed!")