from typing import Sequence, Optional

import numpy as np
import nibabel as nib

from ....torchio import DATA, AFFINE
from ....data.subject import Subject
from .bounds_transform import BoundsTransform, TypeBounds


class Crop(BoundsTransform):
    r"""Crop an image.

    Args:
        cropping: Tuple
            :math:`(w_{ini}, w_{fin}, h_{ini}, h_{fin}, d_{ini}, d_{fin})`
            defining the number of values cropped from the edges of each axis.
            If the initial shape of the image is
            :math:`W \times H \times D`, the final shape will be
            :math:`(- w_{ini} + W - w_{fin}) \times (- h_{ini} + H - h_{fin})
            \times (- d_{ini} + D - d_{fin})`.
            If only three values :math:`(w, h, d)` are provided, then
            :math:`w_{ini} = w_{fin} = w`,
            :math:`h_{ini} = h_{fin} = h` and
            :math:`d_{ini} = d_{fin} = d`.
            If only one value :math:`n` is provided, then
            :math:`w_{ini} = w_{fin} = h_{ini} = h_{fin}
            = d_{ini} = d_{fin} = n`.
    """
    def __init__(
            self,
            cropping: TypeBounds,
            p: float = 1,
            keys: Optional[Sequence[str]] = None,
            ):
        super().__init__(cropping, p=p, keys=keys)
        self.cropping = cropping
        self.args_names = ('cropping',)

    def apply_transform(self, sample) -> Subject:
        low = self.bounds_parameters[::2]
        high = self.bounds_parameters[1::2]
        index_ini = low
        index_fin = np.array(sample.spatial_shape) - high
        for image in self.get_images(sample):
            new_origin = nib.affines.apply_affine(image.affine, index_ini)
            new_affine = image.affine.copy()
            new_affine[:3, 3] = new_origin
            i0, j0, k0 = index_ini
            i1, j1, k1 = index_fin
            image[DATA] = image[DATA][:, i0:i1, j0:j1, k0:k1].clone()
            image[AFFINE] = new_affine
        return sample

    def inverse(self):
        from .pad import Pad
        return Pad(self.cropping)
