import torchio as tio
import numpy as np
import torch


# Helper function for multiprocessing-safe casting
def _to_float32(tensor: torch.Tensor) -> torch.Tensor:
    """Cast TorchIO tensor to float32 (multiprocessing-safe)."""
    return tensor.to(torch.float32)


class PadToMultiple(tio.Transform):
    """
    Pads a 3D image so that each spatial dimension is a multiple of a given number.
    """
    def __init__(self, multiple: int = 16, **kwargs):
        super().__init__(**kwargs)
        self.multiple = multiple

    def apply_transform(self, subject: tio.Subject) -> tio.Subject:
        image_keys = sorted([i for i in subject.keys() if i.startswith("image_")])
        for image_key in image_keys:
            image = subject[image_key]
            shape = image.spatial_shape
            
            target_shape = [int(np.ceil(dim / self.multiple)) * self.multiple for dim in shape]

            pad_transform = tio.CropOrPad(
                target_shape=tuple(target_shape), padding_mode=0
            )
            subject = pad_transform(subject)
            
        return subject 


def pad_fixed_size_collate_fn(batch):
    """
    Pads each image in the batch to the same size and returns a tensor of images and a tensor of labels.
    Assumes that the subject only contains one image and one label.
    """
    
    # Find the maximum spatial dimensions in the batch
    max_shape = np.max([item['image_1'][tio.DATA].shape[1:] for item in batch], axis=0)

    # Make it a multiple of 16
    target_shape = tuple([int(np.ceil(dim / 16)) * 16 for dim in max_shape])
    pad_transform = tio.CropOrPad(
        target_shape=target_shape, padding_mode=0
    )
    
    # Pad each image to the max shape
    for c, item in enumerate(batch):
        image_keys = sorted([i for i in item.keys() if i.startswith("image_")])
        for image_key in image_keys:
            image = item[image_key]
            padded_image = pad_transform(image)
            batch[c][image_key] = padded_image
    
    return batch



def get_image_transforms(do_augmentation: bool = True):
    """
    Get the image transforms for the training and validation sets.
    """
    
    if do_augmentation:
        train_transforms = tio.Compose([
            tio.RandomFlip(axes=('LR', 'AP', 'IS'), flip_probability=0.5),
            tio.RandomAffine(scales=(0.8, 1.2), degrees=45, translation=(15, 15, 15), p=0.3),
            tio.RandomNoise(std=(0.01, 0.10), p=0.2),
            # tio.RandomBiasField(coefficients=0.5, p=0.2),
            # tio.RandomBlur(std=(0, 2), p=0.2),
            tio.RandomGamma(log_gamma=(-0.3, 0.3), p=0.2),
            tio.RandomSwap(patch_size=10, num_iterations=30, p=0.2)
        ])
    else:
        train_transforms = tio.Compose([
            tio.RandomFlip(axes=('LR', 'AP', 'IS'), flip_probability=0.5),
            tio.RandomAffine(scales=(0.8, 1.2), degrees=45, translation=(15, 15, 15), p=0.3),
        ])
    
    # Add type transform to ensure image data loading for validation
    val_transforms = tio.Compose([tio.Lambda(_to_float32)])
    
    return train_transforms, val_transforms