"""Data loader module for loading and preprocessing plant images."""

import os
from typing import Tuple, List, Dict
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import numpy as np


class PhenologyDataset(Dataset):
    """Dataset class for plant phenology images."""
    
    def __init__(
        self,
        csv_file: str,
        root_dir: str,
        transform: transforms.Compose = None,
        split_ratio: Tuple[float, float, float] = (0.7, 0.15, 0.15),
        split: str = 'train',
        random_seed: int = 42
    ):
        """
        Initialize the dataset.
        
        Args:
            csv_file: Path to CSV file with image paths and labels
            root_dir: Root directory containing images
            transform: Optional transform to be applied on images
            split_ratio: Tuple of (train, val, test) ratios
            split: One of 'train', 'val', or 'test'
            random_seed: Random seed for reproducibility
        """
        self.root_dir = root_dir
        self.transform = transform
        self.split = split
        
        # Load CSV
        self.data_frame = pd.read_csv(csv_file)
        
        # Get unique classes and create mapping
        self.classes = sorted(self.data_frame['phase'].unique())
        self.class_to_idx = {cls: idx for idx, cls in enumerate(self.classes)}
        self.idx_to_class = {idx: cls for cls, idx in self.class_to_idx.items()}
        
        # Split data
        np.random.seed(random_seed)
        indices = np.random.permutation(len(self.data_frame))
        
        train_size = int(len(indices) * split_ratio[0])
        val_size = int(len(indices) * split_ratio[1])
        
        if split == 'train':
            self.indices = indices[:train_size]
        elif split == 'val':
            self.indices = indices[train_size:train_size + val_size]
        elif split == 'test':
            self.indices = indices[train_size + val_size:]
        else:
            raise ValueError(f"Invalid split: {split}. Must be 'train', 'val', or 'test'")
        
        self.data_frame = self.data_frame.iloc[self.indices].reset_index(drop=True)
    
    def __len__(self) -> int:
        """Return the size of the dataset."""
        return len(self.data_frame)
    
    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
        """
        Get item at index.
        
        Args:
            idx: Index of item
            
        Returns:
            Tuple of (image tensor, label)
        """
        if torch.is_tensor(idx):
            idx = idx.tolist()
        
        img_name = os.path.join(self.root_dir, self.data_frame.iloc[idx]['image_path'])
        image = Image.open(img_name).convert('RGB')
        label = self.class_to_idx[self.data_frame.iloc[idx]['phase']]
        
        if self.transform:
            image = self.transform(image)
        
        return image, label
    
    def get_class_weights(self) -> torch.Tensor:
        """
        Calculate class weights for handling imbalanced datasets.
        
        Returns:
            Tensor of class weights
        """
        class_counts = self.data_frame['phase'].value_counts()
        total = len(self.data_frame)
        weights = torch.tensor([total / class_counts[cls] for cls in self.classes])
        return weights / weights.sum()


def get_data_loaders(
    csv_file: str,
    root_dir: str,
    batch_size: int = 32,
    num_workers: int = 4,
    split_ratio: Tuple[float, float, float] = (0.7, 0.15, 0.15),
    random_seed: int = 42
) -> Dict[str, DataLoader]:
    """
    Create data loaders for train, validation, and test sets.
    
    Args:
        csv_file: Path to CSV file with image paths and labels
        root_dir: Root directory containing images
        batch_size: Batch size for data loaders
        num_workers: Number of worker processes for data loading
        split_ratio: Tuple of (train, val, test) ratios
        random_seed: Random seed for reproducibility
        
    Returns:
        Dictionary containing train, val, and test data loaders
    """
    # Define transforms
    train_transform = transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(15),
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    val_test_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    # Create datasets
    train_dataset = PhenologyDataset(
        csv_file, root_dir, train_transform, split_ratio, 'train', random_seed
    )
    val_dataset = PhenologyDataset(
        csv_file, root_dir, val_test_transform, split_ratio, 'val', random_seed
    )
    test_dataset = PhenologyDataset(
        csv_file, root_dir, val_test_transform, split_ratio, 'test', random_seed
    )
    
    # Create data loaders
    data_loaders = {
        'train': DataLoader(
            train_dataset, batch_size=batch_size, shuffle=True,
            num_workers=num_workers, pin_memory=True
        ),
        'val': DataLoader(
            val_dataset, batch_size=batch_size, shuffle=False,
            num_workers=num_workers, pin_memory=True
        ),
        'test': DataLoader(
            test_dataset, batch_size=batch_size, shuffle=False,
            num_workers=num_workers, pin_memory=True
        )
    }
    
    return data_loaders, train_dataset.classes, train_dataset.class_to_idx
