from kornia.constants import Resample
from kornia.core import Tensor
from kornia.augmentation._3d.base import RigidAffineAugmentationBase3D
from kornia.augmentation import random_generator as rg
from kornia.geometry import deg2rad, get_affine_matrix3d, warp_affine3d
from kornia.augmentation.random_generator.base import RandomGeneratorBase, UniformDistribution
from kornia.augmentation.utils import _adapted_rsampling, _tuple_range_reader
from kornia.utils.helpers import _extract_device_dtype
import torch
import torch.nn.functional as F

from typing import Any, Dict, Optional, Tuple, Union

# Affine transform
class RandomAffine3DCustom(RigidAffineAugmentationBase3D):
    r"""Apply affine transformation 3D volumes (5D tensor).

    Based on :class:`kornia.augmentation.RandomAffine3D`.

    The transformation is computed so that the center is kept invariant.

    Args:
        degrees: Range of yaw (x-axis), pitch (y-axis), roll (z-axis) to select from.
            If degrees is a number, then yaw, pitch, roll will be generated from the range of (-degrees, +degrees).
            If degrees is a tuple of (min, max), then yaw, pitch, roll will be generated from the range of (min, max).
            If degrees is a list of floats [a, b, c], then yaw, pitch, roll will be generated from (-a, a), (-b, b)
            and (-c, c).
            If degrees is a list of tuple ((a, b), (m, n), (x, y)), then yaw, pitch, roll will be generated from
            (a, b), (m, n) and (x, y).
            Set to 0 to deactivate rotations.
        translate: tuple of maximum absolute fraction for horizontal, vertical and
            depthical translations (dx,dy,dz). For example translate=(a, b, c), then
            horizontal shift will be randomly sampled in the range -img_width * a < dx < img_width * a
            vertical shift will be randomly sampled in the range -img_height * b < dy < img_height * b.
            depthical shift will be randomly sampled in the range -img_depth * c < dz < img_depth * c.
            Will not translate by default.
        scale: scaling factor interval.
            If (a, b) represents isotropic scaling, the scale is randomly sampled from the range a <= scale <= b.
            If ((a, b), (c, d), (e, f)), the scale is randomly sampled from the range a <= scale_x <= b,
            c <= scale_y <= d, e <= scale_z <= f. Will keep original scale by default.
        shears: Range of degrees to select from.
            If shear is a number, a shear to the 6 facets in the range (-shear, +shear) will be applied.
            If shear is a tuple of 2 values, a shear to the 6 facets in the range (shear[0], shear[1]) will be applied.
            If shear is a tuple of 6 values, a shear to the i-th facet in the range (-shear[i], shear[i])
            will be applied.
            If shear is a tuple of 6 tuples, a shear to the i-th facet in the range (-shear[i, 0], shear[i, 1])
            will be applied.
        resample: resample mode from "nearest" (0) or "bilinear" (1).
        same_on_batch: apply the same transformation across the batch.
        align_corners: interpolation flag.
        keepdim: whether to keep the output shape the same as input (True) or broadcast it
          to the batch form (False). Default: False.

    Shape:
        - Input: :math:`(C, D, H, W)` or :math:`(B, C, D, H, W)`, Optional: :math:`(B, 4, 4)`
        - Output: :math:`(B, C, D, H, W)`

    Note:
        Input tensor must be float and normalized into [0, 1] for the best differentiability support.
        Additionally, this function accepts another transformation tensor (:math:`(B, 4, 4)`), then the
        applied transformation will be merged int to the input transformation tensor and returned.

    Examples:
        >>> import torch
        >>> rng = torch.manual_seed(0)
        >>> input = torch.rand(1, 1, 3, 3, 3)
        >>> aug = RandomAffine3D((15., 20., 20.), p=1.)
        >>> aug(input), aug.transform_matrix
        (tensor([[[[[0.4503, 0.4763, 0.1680],
                   [0.2029, 0.4267, 0.3515],
                   [0.3195, 0.5436, 0.3706]],
        <BLANKLINE>
                  [[0.5255, 0.3508, 0.4858],
                   [0.0795, 0.1689, 0.4220],
                   [0.5306, 0.7234, 0.6879]],
        <BLANKLINE>
                  [[0.2971, 0.2746, 0.3471],
                   [0.4924, 0.4960, 0.6460],
                   [0.3187, 0.4556, 0.7596]]]]]), tensor([[[ 0.9722, -0.0603,  0.2262, -0.1381],
                 [ 0.1131,  0.9669, -0.2286,  0.1486],
                 [-0.2049,  0.2478,  0.9469,  0.0102],
                 [ 0.0000,  0.0000,  0.0000,  1.0000]]]))

    To apply the exact augmenation again, you may take the advantage of the previous parameter state:
        >>> input = torch.rand(1, 3, 32, 32, 32)
        >>> aug = RandomAffine3D((15., 20., 20.), p=1.)
        >>> (aug(input) == aug(input, params=aug._params)).all()
        tensor(True)

    """

    def __init__(
        self,
        degrees: Union[
            Tensor,
            float,
            Tuple[float, float],
            Tuple[float, float, float],
            Tuple[Tuple[float, float], Tuple[float, float], Tuple[float, float]],
        ],
        translate: Optional[Union[Tensor, Tuple[float, float, float]]] = None,
        scale: Optional[
            Union[Tensor, Tuple[float, float], Tuple[Tuple[float, float], Tuple[float, float], Tuple[float, float]]]
        ] = None,
        shears: Union[
            None,
            Tensor,
            float,
            Tuple[float, float],
            Tuple[float, float, float, float, float, float],
            Tuple[
                Tuple[float, float],
                Tuple[float, float],
                Tuple[float, float],
                Tuple[float, float],
                Tuple[float, float],
                Tuple[float, float],
            ],
        ] = None,
        resample: Union[str, int, Resample] = Resample.BILINEAR.name,
        same_on_batch: bool = False,
        align_corners: bool = True,
        p: float = 0.5,
        keepdim: bool = True,
    ) -> None:
        super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim)
        self.degrees = degrees
        self.shears = shears
        self.translate = translate
        self.scale = scale

        self.flags = {"resample": Resample.get(resample), "align_corners": align_corners}
        self._param_generator = rg.AffineGenerator3D(degrees, translate, scale, shears)

    def compute_transformation(self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any]) -> Tensor:
        transform: Tensor = get_affine_matrix3d(
            params["translations"],
            params["center"],
            params["scale"],
            params["angles"],
            deg2rad(params["sxy"]),
            deg2rad(params["sxz"]),
            deg2rad(params["syx"]),
            deg2rad(params["syz"]),
            deg2rad(params["szx"]),
            deg2rad(params["szy"]),
        ).to(input)
        return transform

    def apply_transform(
        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
    ) -> Tensor:
        if not isinstance(transform, Tensor):
            raise TypeError(f"Expected the transform to be a Tensor. Gotcha {type(transform)}")

        # Ensure align_corners is a boolean (avoid passing None to affine_grid/grid_sample)
        align = flags.get("align_corners", True)
        if align is None:
            align = True

        return warp_affine3d(
            input,
            transform[:, :3, :],
            (input.shape[-3], input.shape[-2], input.shape[-1]),
            flags["resample"].name.lower(),
            align_corners=bool(align),
        )
    
    def apply_non_transform_mask(
        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
    ) -> Tensor:
        """Process masks corresponding to the inputs that are no transformation applied."""
        return input

    def apply_transform_mask(
        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
    ) -> Tensor:
        """Process masks corresponding to the inputs that are transformed.

        Note:
            Convert "resample" arguments to "nearest" by default.

        """
        resample_method: Optional[Resample]
        if "resample" in flags:
            resample_method = flags["resample"]
            flags["resample"] = Resample.get("nearest")
        output = self.apply_transform(input, params, flags, transform)
        if resample_method is not None:
            flags["resample"] = resample_method
        return output

# Low resolution transform
class RandomLowResTransformGPU(RigidAffineAugmentationBase3D):
    """
    Apply low resolution simulation to 3D volumes (5D tensor).
    """

    def __init__(
        self,
        scale: Tuple[float, float] = (0.3, 1.0),
        crop: Tuple[float, float] = (0.0, 0.0),
        same_on_batch: bool = False,
        resample: Union[str, int, Resample] = Resample.BILINEAR.name,
        p: float = 1.0,
        keepdim: bool = True,
        **kwargs,
    ) -> None:
        super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim)
        self.flags = {"resample": "nearest"} # Use nearest neighbour for now because changing the resampling method for each channel is tricky 
        self._param_generator = ShapeGenerator3D(
            scale=scale,
            crop=crop
        )
    
    def compute_transformation(self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any]) -> Tensor:
        return self.identity_matrix(input)

    @torch.no_grad()
    def apply_transform(
        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
    ) -> Tensor:
        # input shape: (B, C, D, H, W)
        if not isinstance(input, torch.Tensor):
            raise TypeError(f"Expected input to be a Tensor. Got {type(input)}")

        batch_size, C, D, H, W = input.shape

        # params expected to contain 'scale' and 'crop' as tensors of shape (B, 3)
        if params is None or 'scale' not in params or 'crop' not in params:
            raise ValueError("params must contain 'scale' and 'crop' tensors")

        scales = params['scale']  # shape [B, 3]
        crops = params['crop']    # shape [B, 3]

        resample = self.flags.get('resample', 'nearest')

        # Define interpolation modes
        interp_down = resample
        interp_up = resample


        # Process per-channel and per-batch element
        out = input.clone()

        for b in range(batch_size):
            x = input[b]  # [C, D, H, W]

            sx, sy, sz = scales[b]
            # compute downsampled size
            down_D = max(1, int(round(float(sz) * D)))
            down_H = max(1, int(round(float(sy) * H)))
            down_W = max(1, int(round(float(sx) * W)))

            # downsample
            x_down = F.interpolate(
                x.unsqueeze(0),
                size=(down_D, down_H, down_W),
                mode=interp_down,
                align_corners=False if 'linear' in interp_down else None,
            )

            # upsample back to original resolution (keep as 4D tensor [1,1,D,H,W])
            x_up = F.interpolate(
                x_down,
                size=(D, H, W),
                mode=interp_up,
                align_corners=False if 'linear' in interp_up else None,
            ).squeeze(0)  # [C, D, H, W]

            # determine crop fraction and crop size on the upsampled image
            cx, cy, cz = crops[b]
            # interpret crop as fraction of upsampled size to keep
            crop_D = max(1, int(round(float(cz) * D)))
            crop_H = max(1, int(round(float(cy) * H)))
            crop_W = max(1, int(round(float(cx) * W)))

            # choose top-left-front corner within possible range (bias by crop fraction)
            max_z = max(0, D - crop_D)
            max_y = max(0, H - crop_H)
            max_x = max(0, W - crop_W)

            if max_z == 0:
                start_z = 0
            else:
                start_z = int(round(float(cz) * max_z)) if max_z > 0 else 0
            if max_y == 0:
                start_y = 0
            else:
                start_y = int(round(float(cy) * max_y)) if max_y > 0 else 0
            if max_x == 0:
                start_x = 0
            else:
                start_x = int(round(float(cx) * max_x)) if max_x > 0 else 0

            # crop patch from upsampled image (full resolution)
            z1 = start_z
            y1 = start_y
            x1 = start_x
            z2 = z1 + crop_D
            y2 = y1 + crop_H
            x2 = x1 + crop_W

            patch = x_up[:, z1:z2, y1:y2, x1:x2]

            # place patch back into a full-resolution canvas of the original size (zeros elsewhere)
            canvas = torch.zeros((C, D, H, W), dtype=patch.dtype, device=patch.device)
            canvas[:, z1:z2, y1:y2, x1:x2] = patch

            out[b] = canvas

        return out
    
    def apply_non_transform_mask(
        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
    ) -> Tensor:
        """Process masks corresponding to the inputs that are no transformation applied."""
        return input

    def apply_transform_mask(
        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
    ) -> Tensor:
        """Process masks corresponding to the inputs that are transformed.

        Note:
            Convert "resample" arguments to "nearest" by default.

        """
        output = self.apply_transform(input, params, flags, transform)
        return output

class ShapeGenerator3D(RandomGeneratorBase):
    def __init__(
            self, 
            scale: Tuple[float, float],
            crop: Tuple[float, float]
        ) -> None:
        super().__init__()
        self.scale = scale
        self.crop = crop
    
    def make_samplers(self, device: torch.device, dtype: torch.dtype) -> None:
        scale = _tuple_range_reader(self.scale, 3, device, dtype)
        self.scalex_sampler = UniformDistribution(scale[0, 0], scale[0, 1], validate_args=False)
        self.scaley_sampler = UniformDistribution(scale[1, 0], scale[1, 1], validate_args=False)
        self.scalez_sampler = UniformDistribution(scale[2, 0], scale[2, 1], validate_args=False)

        crop = _tuple_range_reader(self.crop, 3, device, dtype)
        self.cropx_sampler = UniformDistribution(crop[0, 0], crop[0, 1], validate_args=False)
        self.cropy_sampler = UniformDistribution(crop[1, 0], crop[1, 1], validate_args=False)
        self.cropz_sampler = UniformDistribution(crop[2, 0], crop[2, 1], validate_args=False)

    def forward(self, batch_shape: Tuple[int, ...], same_on_batch: bool = False) -> Dict[str, torch.Tensor]:
        batch_size = batch_shape[0]

        _device, _dtype = _extract_device_dtype([self.scalex_sampler, self.scaley_sampler, self.scalez_sampler, self.cropx_sampler, self.cropy_sampler, self.cropz_sampler])

        scalex = _adapted_rsampling((batch_size,), self.scalex_sampler, same_on_batch)
        scaley = _adapted_rsampling((batch_size,), self.scaley_sampler, same_on_batch)
        scalez = _adapted_rsampling((batch_size,), self.scalez_sampler, same_on_batch)
        scale = torch.stack([scalex, scaley, scalez], dim=1)

        cropx = _adapted_rsampling((batch_size,), self.cropx_sampler, same_on_batch)
        cropy = _adapted_rsampling((batch_size,), self.cropy_sampler, same_on_batch)
        cropz = _adapted_rsampling((batch_size,), self.cropz_sampler, same_on_batch)
        crop = torch.stack([cropx, cropy, cropz], dim=1)

        return {
            "scale": torch.as_tensor(scale, device=_device, dtype=_dtype),
            "crop": torch.as_tensor(crop, device=_device, dtype=_dtype)
        }

# Flip transforms
class RandomFlipTransformGPU(RigidAffineAugmentationBase3D):
    """
    Apply low resolution simulation to 3D volumes (5D tensor).
    """

    def __init__(
        self,
        flip_axis: int = [0, 1, 2],
        same_on_batch: bool = False,
        p: float = 1.0,
        keepdim: bool = True,
        **kwargs,
    ) -> None:
        super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim)
        # normalize flip_axis into a list of ints
        if isinstance(flip_axis, int):
            self.flip_axis = [flip_axis]
        else:
            self.flip_axis = list(flip_axis)
        
        # generator creates per-batch flip flags for axes (z, y, x)
        self._param_generator = FlipGenerator3D(flip_axis=self.flip_axis)
    
    def compute_transformation(self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any]) -> Tensor:
        return self.identity_matrix(input)

    @torch.no_grad()
    def apply_transform(
        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
    ) -> Tensor:
        
        # input shape: (B, C, D, H, W)
        if not isinstance(input, torch.Tensor):
            raise TypeError(f"Expected input to be a Tensor. Got {type(input)}")

        batch_size, C, D, H, W = input.shape

        # Expect params to contain 'flip' tensor of shape [B, 3] with 0/1 values
        flips = None
        if params is not None and 'flip' in params:
            flips = params['flip']

        out = input.clone()
        # For each batch element, build list of spatial dims to flip (D,H,W -> dims 2,3,4)
        for b in range(batch_size):
            flip_dims = []
            if flips is not None:
                fb = flips[b]
                # fb expected as length-3 tensor for (z,y,x)
                for axis in range(3):
                    if int(fb[axis]) == 1 and axis in self.flip_axis:
                        flip_dims.append(1 + axis)

            if len(flip_dims) > 0:
                out[b] = torch.flip(input[b], dims=tuple(flip_dims))

        return out
    
    def apply_non_transform_mask(
        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
    ) -> Tensor:
        """Process masks corresponding to the inputs that are no transformation applied."""
        return input

    def apply_transform_mask(
        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
    ) -> Tensor:
        """Process masks corresponding to the inputs that are transformed.

        Note:
            Convert "resample" arguments to "nearest" by default.

        """
        output = self.apply_transform(input, params, flags, transform)
        return output

class FlipGenerator3D(RandomGeneratorBase):
    """
    Generate per-batch flip flags for 3 axes (z, y, x).

    Returns a dict with key "flip" and value tensor of shape (B, 3) with 0/1 values.
    Ensures at least one axis is flipped per batch element.
    """
    def __init__(self, flip_axis):
        super().__init__()
        # flip_axis is a list of allowed axes (subset of [0,1,2]).
        if isinstance(flip_axis, int):
            self.flip_axis = [flip_axis]
        else:
            self.flip_axis = list(flip_axis)

    def make_samplers(self, device: torch.device, dtype: torch.dtype) -> None:
        # use uniform samplers per axis and threshold at 0.5
        self._samplers = [UniformDistribution(0.0, 1.0, validate_args=False) for _ in range(3)]

    def forward(self, batch_shape: Tuple[int, ...], same_on_batch: bool = False) -> Dict[str, torch.Tensor]:
        batch_size = batch_shape[0]

        _device, _dtype = _extract_device_dtype(self._samplers)

        samples = []
        for s in self._samplers:
            r = _adapted_rsampling((batch_size,), s, same_on_batch)
            samples.append(r)

        flips = torch.stack(samples, dim=1).to(device=_device, dtype=_dtype)
        flips = (flips > 0.5).to(torch.int8)

        # ensure at least one flip per batch element (choose randomly among allowed axes)
        for b in range(batch_size):
            if flips[b].sum() == 0:
                # pick one allowed axis at random
                if len(self.flip_axis) == 0:
                    # nothing to flip
                    continue
                choice = int(torch.randint(low=0, high=len(self.flip_axis), size=(1,)).item())
                axis = int(self.flip_axis[choice])
                flips[b, axis] = 1

        return {"flip": flips}