from torch.utils.data import DataLoader
import cv2
import numpy as np
import albumentations as A
from albumentations.pytorch import ToTensorV2
from datasets import load_dataset
from torch.utils.data import Dataset
from src.operations import remap_mask

class ClothesDataset(Dataset):
    def __init__(self, dataset, transform = None, label_map = None):
        self.dataset = dataset
        self.transform = transform
        self.label_map = label_map
        
    def __len__(self):
        return len(self.dataset)
    
    def __getitem__(self, idx):
        sample = self.dataset[idx]
        
        # Image
        image = sample["pixel_values"].convert("RGB")
        image = np.array(image, dtype=np.uint8)
        
        # Mask
        mask = np.array(sample["label"], dtype = np.uint8)
        
        # Remap Mask
        mask = remap_mask(mask, self.label_map)
        
        # Apply transform
        if self.transform:
            transformed = self.transform(image=image, mask=mask)
            image = transformed["image"]
            mask = transformed["mask"]
        image = image.contiguous()
        mask = mask.contiguous()
        return image, mask.long()    
        

def get_transforms(config, train=True):
    image_size = config["data"]["image_size"]
    mean = config["data"]["normalization"]["mean"]
    std = config["data"]["normalization"]["std"]
    
    transforms = [
        A.LongestMaxSize(
            max_size = image_size,
            interpolation=cv2.INTER_LINEAR,
            mask_interpolation=cv2.INTER_NEAREST
        ),
        A.PadIfNeeded(
            min_height=image_size,
            min_width=image_size,
            border_mode=cv2.BORDER_CONSTANT,
            fill = 0,
            fill_mask = 0
        )
    ]
    
    if train:
        transforms.append(
            A.HorizontalFlip(
                p = config["data"]["augmentation"]["horizontal_flip_probability"]
            )
        )
    transforms.extend([
        A.Normalize(
            mean = mean,
            std=std,
        ),
        ToTensorV2()
    ])
    
    return A.Compose(transforms)

def create_dataloaders(config):
    # load dataset
    dataset = load_dataset("ckotait/ATRDataset")
    
    # Transforms
    train_transform = get_transforms(config, train=True)
    val_transform = get_transforms(config, train=False)
    
    # Create dataset class
    train_dataset = ClothesDataset(dataset['train'], train_transform, config["segmentation"]["label_map"])
    val_dataset = ClothesDataset(dataset['validation'], val_transform, config["segmentation"]["label_map"])
    test_dataset = ClothesDataset(dataset['test'], val_transform, config["segmentation"]["label_map"])

    # Data loaders
    train_loader = DataLoader(
        train_dataset,
        batch_size=config["data"]["batch_size"],
        shuffle=True,
        num_workers=config["data"]["num_workers"],
        pin_memory=True,
        drop_last=True
    )
    val_loader = DataLoader(
        val_dataset,
        batch_size=config["data"]["batch_size"],
        shuffle=False,
        num_workers=config["data"]["num_workers"],
        pin_memory=True
    )
    test_loader = DataLoader(
        test_dataset,
        batch_size=config["data"]["batch_size"],
        shuffle=False,
        num_workers=config["data"]["num_workers"],
        pin_memory=True
    )

    return train_loader, val_loader, test_loader