class Image:
    '''`Image` is a class to handle medical images with SimpleITK backend and PyTorch tensor support.

    This class provides functionality to work with 2D/3D medical images, handling both regular images
    and segmentation masks. It maintains both SimpleITK and PyTorch representations, and handles
    coordinate transformations between physical, pixel, and normalized coordinate spaces.

    Args:
        itk_image (SimpleITK.Image): The SimpleITK image object
        device (devicetype, optional): Device to store the PyTorch tensors on. Defaults to 'cuda'.
        is_segmentation (bool, optional): Whether the image is a segmentation mask. Defaults to False.
            Set this to True if the image is an integer-valued segmentation mask.
            This is used to convert integer labels to one-hot encoding.
            See `max_seg_label`, `background_seg_label`, `seg_preprocessor` for more details on how to manipulate integer label images.
        max_seg_label (int, optional): Maximum label value for segmentation. Values above this are set to background_seg_label.
            Set to None by default, meaning no label clipping is done.
        background_seg_label (int, optional): Label value representing background in segmentations. Defaults to -1 (background is included as a one-hot label)
        is_onehot (bool, optional): If True, use one-hot interpolation (bilinear/trilinear on one-hot channels). If False and FFO is available, use generic-label fused sampler. Defaults to False. When FFO is not available, this is forced to True.
        seg_preprocessor (callable, optional): Function to preprocess segmentation arrays. Defaults to identity function.
        orientation (str, optional): Reorient the image to this orientation. Defaults to None.
        spacing (array-like, optional): Custom spacing for the image. If None, uses SimpleITK values.
        direction (array-like, optional): Custom direction matrix. If None, uses SimpleITK values.
        origin (array-like, optional): Custom origin point. If None, uses SimpleITK values.
        center (array-like, optional): Custom center point to recalibrate origin. If provided, origin is recalculated.

    Attributes:
        array (torch.Tensor): PyTorch tensor representation of the image
            This is of size [1, C, *dims] where C is the number of channels / labels
        itk_image (SimpleITK.Image): Original SimpleITK image
        channels (int): Number of channels in the image
        dims (int): Dimensionality of the image (2 or 3)
        torch2phy (torch.Tensor): Transform matrix from normalized to physical coordinates
        phy2torch (torch.Tensor): Transform matrix from physical to normalized coordinates
        device (devicetype): Device where PyTorch tensors are stored
    '''
    def __init__(self, itk_image: sitk.SimpleITK.Image,
                 device: devicetype = 'cuda',
                 dtype: torch.dtype = None,
                 is_segmentation=False, max_seg_label=None,
                 background_seg_label=-1, is_onehot: bool = False, seg_preprocessor=lambda x: x,
                 orientation: str = None,
                 winsorize: bool = False,
                 winsorize_percentile: Tuple[float, float] = (1.0, 99.0),
                spacing=None, direction=None, origin=None, center=None) -> None:

        if orientation is not None:
            itk_image = sitk.DICOMOrient(itk_image, orientation)

        self.itk_image = itk_image
        if dtype is None:
            dtype = torch.float32
        # check for segmentation parameters
        # if `is_segmentation` is False, then just treat this as an image with given dtype
        if not is_segmentation:
            self.interpolation_mode = None
            self.is_onehot = False
            itk_img = sitk.GetArrayFromImage(itk_image).astype(float)
            if winsorize:
                itk_img = winsorize_image(itk_img, winsorize_percentile[0], winsorize_percentile[1])
            self.array = torch.from_numpy(itk_img).to(device, dtype)
            channels = itk_image.GetNumberOfComponentsPerPixel()
            self.channels = channels
            # assert channels == 1, "Only single channel images supported"
            if channels > 1:
                logger.warning("Image has multiple channels, make sure its not a spatial dimension")
                # permute the channel dimension to the front
                ndim = self.array.ndim
                self.array = self.array.permute([ndim-1] + list(range(ndim-1))).contiguous() # permute to [C, H, W, D] from [H, W, D, C]
            else:
                self.array.unsqueeze_(0)
            # add batch dimension
            self.array.unsqueeze_(0)
        else:
            # When FFO is not available, force one-hot mode so generic-label path is never used
            if not fi.use_ffo:
                is_onehot = True
            self.is_onehot = is_onehot
            array = torch.from_numpy(sitk.GetArrayFromImage(itk_image).astype(int)).float().to(device)
            # preprocess segmentation if provided by user
            array = seg_preprocessor(array)
            if max_seg_label is not None:
                array[array > max_seg_label] = background_seg_label if is_onehot else array.min()
            # convert to one-hot encoding
            if is_onehot:
                array = integer_to_onehot(array.long(), background_label=background_seg_label, max_label=max_seg_label, dtype=dtype)[None]  # [1, C, H, W, D]
            else:
                array = array[None, None]  # [1, 1, H, W, D]
            # set these values
            self.array = array
            self.channels = array.shape[1]
            self.interpolation_mode = 'genericlabel' if (not is_onehot) else None
        # initialize matrix for pixel to physical
        dims = itk_image.GetDimension()
        self.dims = dims
        if dims not in [2, 3]:
            raise NotImplementedError("Image class only supports 2D/3D images.")

        # custom spacing if not provided use simpleitk values
        spacing = np.array(itk_image.GetSpacing())[None] if spacing is None else np.array(spacing)[None]
        origin = np.array(itk_image.GetOrigin())[None] if origin is None else np.array(origin)[None]
        direction = np.array(itk_image.GetDirection()).reshape(dims, dims) if direction is None else np.array(direction).reshape(dims, dims)
        if center is not None:
            print("Center location provided, recalibrating origin.")
            origin = center - np.matmul(direction, ((np.array(itk_image.GetSize())*spacing/2).squeeze())[:, None]).T

        px2phy = np.eye(dims+1)
        px2phy[:dims, -1] = origin
        px2phy[:dims, :dims] = direction
        px2phy[:dims, :dims] = px2phy[:dims, :dims] * spacing
        # generate mapping from torch to px
        torch2px = np.eye(dims+1)
        scaleterm = (np.array(itk_image.GetSize())-1)*0.5
        torch2px[:dims, :dims] = np.diag(scaleterm)
        torch2px[:dims, -1] = scaleterm
        # save the mapping from physical to torch and vice versa
        self.torch2phy = torch.from_numpy(np.matmul(px2phy, torch2px)).to(device).float().unsqueeze_(0)
        self.phy2torch = torch.inverse(self.torch2phy[0]).float().unsqueeze_(0)
        # also save intermediates just in case (as numpy arrays)
        self._torch2px = torch2px
        self._px2phy = px2phy
        # keep these as well
        self.torch2px = torch.from_numpy(self._torch2px).to(device).float().unsqueeze_(0)
        self.px2torch = torch.inverse(self.torch2px[0]).float().unsqueeze_(0)
        self.px2phy = torch.from_numpy(self._px2phy).to(device).float().unsqueeze_(0)
        self.phy2px = torch.inverse(self.px2phy[0]).float().unsqueeze_(0)

        self._px2phy = px2phy
        self.device = device

    @classmethod
    def load_file(cls, image_path:str, *args: Any, **kwargs: Any) -> 'Image':
        '''Load an image from a file.

        Args:
            image_path (str): Path to the image file
            *args: Additional arguments to pass to the Image constructor
            **kwargs: Additional arguments to pass to the Image constructor

        Returns:
            Image: An instance of the Image class
        '''
        itk_image = sitk.ReadImage(image_path)
        return cls(itk_image, *args, **kwargs)

    @property
    def shape(self) -> Union[torch.Size, List, Tuple]:
        '''Get the shape of the image.

        Returns:
            torch.Size: Shape of the image
        '''
        return self.array.shape

    @property
    def is_array_present(self) -> bool:
        '''Check if the PyTorch tensor representation of the image is present.
        This is needed because the BatchedImages may delete the array of the image.

        Returns:
            bool: True if the PyTorch tensor representation of the image is present, False otherwise
        '''
        return hasattr(self, 'array')

    def delete_array(self):
        '''Delete the PyTorch tensor representation of the image.

        This is a placeholder function to be replaced with batched images.
        '''
        if self.is_array_present:
            del self.array

    def concat(self, *others, optimize_memory: bool = True):
        ''' alias for concatenate '''
        return self.concatenate(*others, optimize_memory=optimize_memory)

    def concatenate(self, *others, optimize_memory: bool = True):
        '''
        others is a list of Images or list of lists (in which case the user accidentally passed a list of images)

            optimize_memory: if True, delete the arrays of the other images after concatenation

        Example:
            t1 = Image.load_file(t1_path)
            t2 = Image.load_file(t2_path)
            flair = Image.load_file(flair_path)
            t1.concatenate(t2, flair, optimize_memory=True)   # deletes the arrays of t2 and flair after concatenation
        '''
        if len(others) == 0:
            return self

        check_and_raise_cond(self.is_array_present, "Image must have a PyTorch tensor representation to concatenate", ValueError)
        if isinstance(others[0], list) and len(others) == 1:
            others = others[0]
        else:
            check_and_raise_cond(all([isinstance(other, Image) for other in others]), "All elements of others must be of type Image", TypeError)
        # check if all images have the same shape
        shapes = [x.array.shape[2:] for x in others]
        base_shape = self.array.shape[2:]
        check_and_raise_cond(all([x == base_shape for x in shapes]), "All images must have the same shape", ValueError)
        check_and_raise_cond(all([x.is_array_present for x in others]), "All images must have a PyTorch tensor representation", ValueError)
        check_and_raise_cond(all([self.array.device == other.array.device for other in others]), "Images must be on the same device", ValueError)
        check_and_raise_cond(all([torch.allclose(self.phy2torch, other.phy2torch) for other in others]), "Images reside in different physical spaces, using the first image's physical space", logger.warning)

        self.array = torch.cat([self.array] + [other.array for other in others], dim=1)
        self.channels = self.array.shape[1]

        if optimize_memory:
            logger.debug("Deleting the arrays of the other images after concatenation")
            for other in others:
                other.delete_array()
        return self

    def to(self, device_or_dtype: Union[devicetype, torch.dtype]):
        '''
        Move the image to a different device or change its dtype.

        Args:
            device_or_dtype (Union[devicetype, torch.dtype]): Device to move the image to,
                or dtype to convert the image to

        Returns:
            Image: The image on the new device or with the new dtype
        '''
        self.array = self.array.to(device_or_dtype)
        if isinstance(device_or_dtype, (str, torch.device)):
            self.device = device_or_dtype
        return self

    def __del__(self):
        '''Delete the SimpleITK image and all intermediate variables.'''
        del self.itk_image
        if self.is_array_present:
            del self.array
        del self.torch2phy
        del self.phy2torch
        del self._torch2px
        del self._px2phy
        del self.torch2px
        del self.px2torch
        del self.px2phy
        del self.phy2px

class BatchedImages:
    '''A class to handle batches of Image objects efficiently.

    This class provides functionality to work with multiple Image objects as a batch,
    with options for memory optimization and broadcasting. All images in a batch must
    have the same shape.

    Args:
        images (Union[Image, List[Image]]): Single Image object or list of Image objects
        optimize_memory (bool, optional): Flag for memory optimization (reserved for future use)

    Attributes:
        images (List[Image]): List of Image objects in the batch
        n_images (int): Number of images in the batch
        interpolate_mode (str): Interpolation mode ('bilinear' for 2D, 'trilinear' for 3D)
        broadcasted (bool): Whether the batch has been broadcasted
        device (devicetype): Device where the image tensors are stored
        dims (int): Dimensionality of the images (2 or 3)

    Raises:
        ValueError: If batch is empty or images have different shapes
        TypeError: If any element is not an Image object
    '''
    def __init__(self, images: Union[Image, List[Image]], optimize_memory: bool = None) -> None:
        if isinstance(images, Image):
            images = [images]
        self.images = images

        if len(self.images) == 0:
            raise ValueError("BatchedImages must have at least one image")
        # check if all images have a PyTorch tensor representation
        # Check if all images are of type Image
        # check if all images have the same shape
        check_and_raise_cond(all([image.is_array_present for image in self.images]), "All images must have a PyTorch tensor representation", ValueError)
        check_and_raise_cond(all([isinstance(image, Image) for image in self.images]), "All images must be of type Image", TypeError)
        shapes = [x.array.shape[2:] for x in self.images]
        check_and_raise_cond(all([x == shapes[0] for x in shapes]), "All images must have the same shape", ValueError)

        # set the number of images
        self.n_images = len(self.images)
        # derive the interpolation mode from the first image if it is present, otherwise use bilinear/trilinear
        first_interp = getattr(self.images[0], 'interpolation_mode', None)
        if first_interp is not None:
            self.interpolate_mode = first_interp
        else:
            self.interpolate_mode = 'bilinear' if len(self.images[0].shape) == 4 else 'trilinear'
        self.broadcasted = False
        # create a batch tensor
        self.batch_tensor = torch.cat([x.array for x in self.images], dim=0) if self.n_images > 1 else self.images[0].array
        # delete the arrays of the images if multiple images
        if optimize_memory and len(self.images) > 1:
            logger.info("Deleting the arrays of the images")
            for image in self.images:
                image.delete_array()
        # create metadata
        self.torch2phy = torch.cat([x.torch2phy for x in self.images], dim=0)
        self.phy2torch = torch.cat([x.phy2torch for x in self.images], dim=0)
        self.torch2px = torch.cat([x.torch2px for x in self.images], dim=0)
        self.px2torch = torch.cat([x.px2torch for x in self.images], dim=0)
        self.px2phy = torch.cat([x.px2phy for x in self.images], dim=0)
        self.phy2px = torch.cat([x.phy2px for x in self.images], dim=0)
        # sharding info
        self.is_sharded = False

    def set_device(self, device_name):
        # set the device for the batch tensor
        self.batch_tensor = self.batch_tensor.to(device_name)
        self.torch2phy = self.torch2phy.to(device_name)
        self.phy2torch = self.phy2torch.to(device_name)
        self.torch2px = self.torch2px.to(device_name)
        self.px2torch = self.px2torch.to(device_name)
        self.px2phy = self.px2phy.to(device_name)
        self.phy2px = self.phy2px.to(device_name)

    def to(self, device_or_dtype: Union[devicetype, torch.dtype]):
        '''
        Move the batch to a different device or change its dtype.
        '''
        self.batch_tensor = self.batch_tensor.to(device_or_dtype)
        return self

    def _shard_dim(self, dim_to_shard, rank=None):
        if self.is_sharded:
            logger.warning("Batch is already sharded, cannot shard again")
            return
        if rank is None:
            rank = parallel_state.get_parallel_state().get_rank()
        gp_group = parallel_state.get_parallel_state().get_current_gp_group()
        gp_size = len(gp_group)
        size = self.batch_tensor.shape[dim_to_shard+2]
        chunk_sizes = divide_size_into_chunks(size, gp_size)
        local_rank = gp_group.index(rank)
        # check
        start = sum(chunk_sizes[:local_rank])
        end = sum(chunk_sizes[:local_rank+1])
        # store the full tensor into batch_tensor_full, and sharded version in unshareded
        self.batch_tensor_full = self.batch_tensor.clone()
        self.batch_tensor = self.batch_tensor_full.narrow_copy(dim_to_shard+2, start, end-start).clone()
        print(f"Sharded batch tensor shape: {self.batch_tensor.shape} and rank: {local_rank}/{gp_size}, rank: {rank}")
        self._shard_start = start
        self._shard_end = end
        self._dim_to_shard = dim_to_shard
        self.is_sharded = True

    def _save_shards(self, dim_to_shard):
        gp_group = parallel_state.get_parallel_state().get_current_gp_group()
        gp_size = len(gp_group)
        # get size and chunk size
        size = self.batch_tensor.shape[dim_to_shard+2]
        chunk_size = (size + gp_size - 1) // gp_size
        chunks = []
        for _ in range(gp_size):
            start = _ * chunk_size
            end = min(start + chunk_size, size)
            chunks.append((start, end))
        self._sharded_chunks = chunks


    def __call__(self):
        '''Get the batch of images.
        '''
        if self.broadcasted:
            minusones = [-1] * (len(self.images[0].shape) - 1)
            return self.batch_tensor.expand(self.n_images, *minusones)
        else:
            return self.batch_tensor

    def get_interpolator_type(self):
        '''
        get the interpolator type to be used in `fireants_interpolator`

        - since the grid sampler only accepts only bilinear / nearest, we need to return 'bilinear' for bilinear/trilinear
        '''
        if self.interpolate_mode in ['bilinear', 'trilinear']:
            return 'bilinear'
        return self.interpolate_mode

    def broadcast(self, n):
        '''Broadcast the batch to n channels.

        Args:
            n (int): Number of channels to broadcast to

        Raises:
            ValueError: If batch size is not 1
        '''
        if not self.broadcasted and self.n_images != 1:
            raise ValueError("Batch size must be 1 to broadcast")
        self.broadcasted = True
        self.n_images = n

    def __del__(self):
        '''Delete all Image objects in the batch.'''
        del self.batch_tensor
        del self.torch2phy
        del self.phy2torch

    @property
    def device(self):
        '''Get the device where the image tensors are stored.'''
        return self.batch_tensor.device

    @property
    def dims(self):
        '''Get the dimensionality of the images.'''
        return self.images[0].dims

    def size(self):
        '''Get the number of images in the batch.'''
        return self.n_images

    @property
    def shape(self):
        '''Get the shape of the images.'''
        # shape = list(self.images[0].shape)
        shape = list(self.batch_tensor.shape)
        shape[0] = self.n_images
        return shape

    def get_torch2phy(self):
        return self.torch2phy

    def get_phy2torch(self):
        return self.phy2torch