import os
from torch.utils.data import DataLoader
from mindglow.seg2d.dataset.segdataset import SegDataset
from torchvision import transforms as T
import torch

class ResizeTransform:
    """Resize image and mask to a given size."""
    def __init__(self, size):
        self.size = size if isinstance(size, tuple) else (size, size)

    def __call__(self, image, mask):
        # image: (C, H, W) float tensor, mask: (H, W) long tensor
        from torchvision.transforms import functional as F
        size = list(self.size)  # convert tuple to list for type checker
        # Resize image (bilinear) and mask (nearest to keep class indices)
        image = F.resize(image.unsqueeze(0), size, interpolation=T.InterpolationMode.BILINEAR).squeeze(0)
        mask = F.resize(mask.unsqueeze(0).float(), size, interpolation=T.InterpolationMode.NEAREST).squeeze(0).long()
        return image, mask


def create_dataloaders(config):
    """
    Create train and validation DataLoaders from a configuration dictionary.

    The config must contain a 'dataset' key. It can be:
        - a string: path to a folder with train/ and val/ subfolders (each containing 'images/' and 'masks/')
        - a dict: parameters passed directly to SegDataset (e.g., format, image_suffix, etc.)

    Additional keys used:
        - batch_size (int)
        - num_workers (int, default 4)
        - image_size (int or tuple, optional): if provided, images and masks are resized to this size.
        - image_suffix (str, default '.png'): used only when dataset is a string.
        - mask_suffix (str, optional): used only when dataset is a string; defaults to image_suffix.

    Returns:
        train_loader, val_loader
    """
    dataset_cfg = config.get('dataset')
    if dataset_cfg is None:
        raise ValueError("Config must contain a 'dataset' key")

    batch_size = config['batch_size']
    num_workers = config.get('num_workers', 4)
    pin_memory = True

    # Prepare a resize transform if image_size is provided
    image_size = config.get('image_size')
    resize_transform = ResizeTransform(image_size) if image_size else None

    def build_dataset(split):
        # If dataset_cfg is a string, assume a simple folder structure
        if isinstance(dataset_cfg, str):
            return SegDataset(
                config=dataset_cfg,
                format='folder',
                split=split,
                image_suffix=config.get('image_suffix', '.png'),
                mask_suffix=config.get('mask_suffix', None),
                transforms=resize_transform
            )
        else:
            # dataset_cfg is a dict: copy and add split, and possibly resize transform
            ds_kwargs = dataset_cfg.copy()
            ds_kwargs['split'] = split

            # If we have a resize transform, we need to combine it with any existing transforms
            if resize_transform is not None:
                existing = ds_kwargs.get('transforms')
                if existing is None:
                    ds_kwargs['transforms'] = resize_transform
                else:
                    from torchvision.transforms import Compose
                    # Convert existing to a list of transforms
                    if isinstance(existing, Compose):
                        existing = existing.transforms   # Compose stores transforms in a list
                    elif not isinstance(existing, (list, tuple)):
                        existing = [existing]
                    elif isinstance(existing, tuple):
                        existing = list(existing)
                    # Now existing is definitely a list
                    ds_kwargs['transforms'] = Compose([resize_transform] + existing)
            return SegDataset(**ds_kwargs)

    train_dataset = build_dataset('train')
    val_dataset = build_dataset('val')

    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=pin_memory,
        drop_last=True,
    )
    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=pin_memory,
    )
    return train_loader, val_loader