# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_vision_core.ipynb.

# %% auto 0
__all__ = ['med_img_reader', 'MetaResolver', 'MedBase', 'MedImage', 'MedMask']

# %% ../nbs/01_vision_core.ipynb 2
from .vision_plot import *
from fastai.data.all import *
from torchio import ScalarImage, LabelMap, ToCanonical, Resample
import copy

# %% ../nbs/01_vision_core.ipynb 5
def _preprocess(obj, apply_reorder, target_spacing):
    """
    Preprocesses the given object.

    Args:
        obj: The object to preprocess.
        apply_reorder: Whether to reorder the object.
        target_spacing: Whether to resample the object.

    Returns:
        The preprocessed object and its original size.
    """
    if apply_reorder:
        transform = ToCanonical()
        obj = transform(obj)

    original_size = obj.shape[1:]

    if target_spacing and not all(np.isclose(obj.spacing, target_spacing)):
        transform = Resample(target_spacing)
        obj = transform(obj)

    if MedBase.affine_matrix is None:
        MedBase.affine_matrix = obj.affine

    return obj, original_size

# %% ../nbs/01_vision_core.ipynb 6
def _load_and_preprocess(file_path, apply_reorder, target_spacing, dtype):
    """
    Helper function to load and preprocess an image.

    Args:
        file_path: Image file path.
        apply_reorder: Whether to reorder data for canonical (RAS+) orientation.
        target_spacing: Whether to resample image to different voxel sizes and dimensions.
        dtype: Desired datatype for output.

    Returns:
        tuple: Original image, preprocessed image, and its original size.
    """
    org_img = LabelMap(file_path) if dtype is MedMask else ScalarImage(file_path) #_load(file_path, dtype=dtype) 
    input_img, org_size = _preprocess(org_img, apply_reorder, target_spacing)
    
    return org_img, input_img, org_size

# %% ../nbs/01_vision_core.ipynb 7
def _multi_channel(image_paths: L | list, apply_reorder: bool, target_spacing: list, only_tensor: bool, dtype):
    """
    Load and preprocess multisequence data.

    Args:
        image_paths: List of image paths (e.g., T1, T2, T1CE, DWI).
        apply_reorder: Whether to reorder data for canonical (RAS+) orientation.
        target_spacing: Whether to resample image to different voxel sizes and dimensions.
        only_tensor: Whether to return only image tensor.
        dtype: Desired datatype for output.

    Returns:
        torch.Tensor: A stacked 4D tensor, if `only_tensor` is True.
        tuple: Original image, preprocessed image, original size, if `only_tensor` is False.
    """
    image_data = [_load_and_preprocess(image, apply_reorder, target_spacing, dtype) for image in image_paths]
    org_img, input_img, org_size = image_data[-1]

    tensor = torch.stack([img.data[0] for _, img, _ in image_data], dim=0)
    
    if only_tensor: 
       return dtype(tensor) 

    input_img.set_data(tensor)
    return org_img, input_img, org_size

# %% ../nbs/01_vision_core.ipynb 8
def med_img_reader(file_path: str | Path | L | list, apply_reorder: bool = False, target_spacing: list = None, 
                   only_tensor: bool = True, dtype = torch.Tensor):
    """Loads and preprocesses a medical image.

    Args:
        file_path: Path to the image. Can be a string, Path object or a list.
        apply_reorder: Whether to reorder the data to be closest to canonical 
            (RAS+) orientation. Defaults to False.
        target_spacing: Whether to resample image to different voxel sizes and 
            image dimensions. Defaults to None.
        only_tensor: Whether to return only image tensor. Defaults to True.
        dtype: Datatype for the return value. Defaults to torch.Tensor.

    Returns:
        The preprocessed image. Returns only the image tensor if 
        only_tensor is True, otherwise returns original image, 
        preprocessed image, and original size.
    """
    
    if isinstance(file_path, (list, L)):
        return _multi_channel(file_path, apply_reorder, target_spacing, only_tensor, dtype)

    org_img, input_img, org_size = _load_and_preprocess(
        file_path, apply_reorder, target_spacing, dtype)

    if only_tensor:
        return dtype(input_img.data.type(torch.float))

    return org_img, input_img, org_size

# %% ../nbs/01_vision_core.ipynb 10
class MetaResolver(type(torch.Tensor), metaclass=BypassNewMeta):
    """
    A class to bypass metaclass conflict:
    https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/data/batch.html
    """
    pass

# %% ../nbs/01_vision_core.ipynb 11
class MedBase(torch.Tensor, metaclass=MetaResolver):
    """A class that represents an image object.
    Metaclass casts `x` to this class if it is of type `cls._bypass_type`."""
    
    _bypass_type = torch.Tensor
    _show_args = {'cmap':'gray'}
    target_spacing, apply_reorder = None, False
    affine_matrix = None

    @classmethod
    def create(cls, fn: (Path, str, L, list, torch.Tensor), **kwargs) -> torch.Tensor: 
        """
        Opens a medical image and casts it to MedBase object.
        If `fn` is a torch.Tensor, it's cast to MedBase object.

        Args:
            fn : (Path, str, torch.Tensor)
                Image path or a 4D torch.Tensor.
            kwargs : dict
                Additional parameters for the medical image reader.

        Returns:
            torch.Tensor : A 4D tensor as a MedBase object.
        """
        if isinstance(fn, torch.Tensor):
            return cls(fn)

        return med_img_reader(fn, target_spacing=cls.target_spacing, apply_reorder=cls.apply_reorder, dtype=cls)

    def __new__(cls, x, **kwargs):
        """Creates a new instance of MedBase from a tensor."""
        if isinstance(x, torch.Tensor):
            # Create tensor of the same type and copy data
            res = torch.Tensor._make_subclass(cls, x.data, x.requires_grad)
            # Copy any additional attributes
            if hasattr(x, 'affine_matrix'):
                res.affine_matrix = x.affine_matrix
            return res
        else:
            # Handle other types by converting to tensor first
            tensor = torch.as_tensor(x, **kwargs)
            return cls.__new__(cls, tensor)

    def new_empty(self, size, **kwargs):
        """Create a new empty tensor of the same type."""
        # Create new tensor with same type and device/dtype
        kwargs.setdefault('dtype', self.dtype)
        kwargs.setdefault('device', self.device)
        new_tensor = torch.empty(size, **kwargs)
        # Use __new__ to create proper subclass instance
        return self.__class__.__new__(self.__class__, new_tensor)

    def __copy__(self):
        """Shallow copy implementation."""
        copied = self.__class__.__new__(self.__class__, self.clone())
        # Copy class attributes
        if hasattr(self, 'affine_matrix'):
            copied.affine_matrix = self.affine_matrix
        return copied

    def __deepcopy__(self, memo):
        """Deep copy implementation."""
        # Create a deep copy of the tensor data
        copied_data = self.clone()
        copied = self.__class__.__new__(self.__class__, copied_data)
        # Deep copy class attributes
        if hasattr(self, 'affine_matrix') and self.affine_matrix is not None:
            copied.affine_matrix = copy.deepcopy(self.affine_matrix, memo)
        else:
            copied.affine_matrix = None
        return copied

    @classmethod
    def item_preprocessing(cls, target_spacing: (list, int, tuple), apply_reorder: bool):
        """
        Changes the values for the class variables `target_spacing` and `apply_reorder`.

        Args:
            target_spacing : (list, int, tuple)
                A list with voxel spacing.
            apply_reorder : bool
                Whether to reorder the data to be closest to canonical (RAS+) orientation.
        """
        cls.target_spacing = target_spacing
        cls.apply_reorder = apply_reorder

    def show(self, ctx=None, channel: int = 0, slice_index: int = None, anatomical_plane: int = 0, **kwargs):
        """
        Displays the Medimage using `merge(self._show_args, kwargs)`.

        Args:
            ctx : Any, optional
                Context to use for the display. Defaults to None.
            channel : int, optional
                The channel of the image to be displayed. Defaults to 0.
            slice_index : int or None, optional
                Index of the images to be displayed. Defaults to None.
            anatomical_plane : int, optional
                Anatomical plane of the image to be displayed. Defaults to 0.
            kwargs : dict, optional
                Additional parameters for the show function.

        Returns:
            Shown image.
        """
        return show_med_img(
            self, ctx=ctx, channel=channel, slice_index=slice_index, 
            anatomical_plane=anatomical_plane, voxel_size=self.target_spacing,  
            **merge(self._show_args, kwargs)
        )

    def __repr__(self) -> str:
        """Returns the string representation of the MedBase instance."""
        return f'{self.__class__.__name__} mode={self.mode} size={"x".join([str(d) for d in self.size])}'

# %% ../nbs/01_vision_core.ipynb 12
class MedImage(MedBase):
    """Subclass of MedBase that represents an image object."""
    pass

# %% ../nbs/01_vision_core.ipynb 13
class MedMask(MedBase):
    """Subclass of MedBase that represents an mask object."""
    _show_args = {'alpha':0.5, 'cmap':'tab20'}
