"""
Custom Dataset Classes - PyTorch Neural Networks
Project: PyTorch Neural Networks
Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
Description: Custom dataset classes for various data types
"""

import torch
from torch.utils.data import Dataset
import numpy as np
from PIL import Image
import os


class ImageDataset(Dataset):
    """
    Custom image dataset class
    
    Project: PyTorch Neural Networks
    Author: RSK World
    Website: https://rskworld.in
    """
    
    def __init__(self, image_paths, labels, transform=None):
        """
        Initialize image dataset
        
        Args:
            image_paths: List of image file paths
            labels: List of labels
            transform: Optional transform to apply
        """
        self.image_paths = image_paths
        self.labels = labels
        self.transform = transform
    
    def __len__(self):
        return len(self.image_paths)
    
    def __getitem__(self, idx):
        image_path = self.image_paths[idx]
        label = self.labels[idx]
        
        # Load image
        image = Image.open(image_path).convert('RGB')
        
        if self.transform:
            image = self.transform(image)
        
        return image, label


class SequenceDataset(Dataset):
    """
    Custom sequence dataset class
    
    Project: PyTorch Neural Networks
    Author: RSK World
    Website: https://rskworld.in
    """
    
    def __init__(self, sequences, labels, transform=None):
        """
        Initialize sequence dataset
        
        Args:
            sequences: Array of sequences (N, seq_len, features)
            labels: Array of labels
            transform: Optional transform to apply
        """
        self.sequences = torch.FloatTensor(sequences)
        self.labels = torch.LongTensor(labels)
        self.transform = transform
    
    def __len__(self):
        return len(self.sequences)
    
    def __getitem__(self, idx):
        sequence = self.sequences[idx]
        label = self.labels[idx]
        
        if self.transform:
            sequence = self.transform(sequence)
        
        return sequence, label


class TabularDataset(Dataset):
    """
    Custom tabular dataset class
    
    Project: PyTorch Neural Networks
    Author: RSK World
    Website: https://rskworld.in
    """
    
    def __init__(self, features, labels, transform=None):
        """
        Initialize tabular dataset
        
        Args:
            features: Array of features (N, features)
            labels: Array of labels
            transform: Optional transform to apply
        """
        self.features = torch.FloatTensor(features)
        self.labels = torch.LongTensor(labels)
        self.transform = transform
    
    def __len__(self):
        return len(self.features)
    
    def __getitem__(self, idx):
        feature = self.features[idx]
        label = self.labels[idx]
        
        if self.transform:
            feature = self.transform(feature)
        
        return feature, label