from collections.abc import Sequence

import numpy as np
import SimpleITK as sitk
import torch

from ....data.io import nib_to_sitk
from ....data.subject import Subject
from ...fourier import FourierTransform
from ...intensity_transform import IntensityTransform
from .. import RandomTransform


class RandomMotion(RandomTransform, IntensityTransform, FourierTransform):
    r"""Add random MRI motion artifact.

    Magnetic resonance images suffer from motion artifacts when the subject
    moves during image acquisition. This transform follows
    [Shaw et al., 2019 ](http://proceedings.mlr.press/v102/shaw19a.html) to
    simulate motion artifacts for data augmentation.

    Args:
        degrees: Tuple $(a, b)$ defining the rotation range in degrees of
            the simulated movements. The rotation angles around each axis are
            $(\theta_1, \theta_2, \theta_3)$,
            where $\theta_i \sim \mathcal{U}(a, b)$.
            If only one value $d$ is provided,
            $\theta_i \sim \mathcal{U}(-d, d)$.
            Larger values generate more distorted images.
        translation: Tuple $(a, b)$ defining the translation in mm of
            the simulated movements. The translations along each axis are
            $(t_1, t_2, t_3)$,
            where $t_i \sim \mathcal{U}(a, b)$.
            If only one value $t$ is provided,
            $t_i \sim \mathcal{U}(-t, t)$.
            Larger values generate more distorted images.
        num_transforms: Number of simulated movements.
            Larger values generate more distorted images.
        image_interpolation: See Interpolation.
        **kwargs: See [`Transform`][torchio.transforms.Transform] for additional
            keyword arguments.

    Warning:
        Large numbers of movements lead to longer execution times for
        3D images.
    """

    def __init__(
        self,
        degrees: float | tuple[float, float] = 10,
        translation: float | tuple[float, float] = 10,  # in mm
        num_transforms: int = 2,
        image_interpolation: str = 'linear',
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.degrees_range = self.parse_degrees(degrees)
        self.translation_range = self.parse_translation(translation)
        if num_transforms < 1 or not isinstance(num_transforms, int):
            message = (
                'Number of transforms must be a strictly positive natural'
                f'number, not {num_transforms}'
            )
            raise ValueError(message)
        self.num_transforms = num_transforms
        self.image_interpolation = self.parse_interpolation(
            image_interpolation,
        )

    def apply_transform(self, subject: Subject) -> Subject:
        images_dict = self.get_images_dict(subject)
        if not images_dict:
            return subject

        times_by_name: dict[str, np.ndarray] = {}
        degrees_by_name: dict[str, np.ndarray] = {}
        translation_by_name: dict[str, np.ndarray] = {}
        interpolation_by_name: dict[str, str] = {}
        for name, image in images_dict.items():
            params = self.get_params(
                self.degrees_range,
                self.translation_range,
                self.num_transforms,
                is_2d=image.is_2d(),
            )
            times_params, degrees_params, translation_params = params
            times_by_name[name] = times_params
            degrees_by_name[name] = degrees_params
            translation_by_name[name] = translation_params
            interpolation_by_name[name] = self.image_interpolation
        transform = Motion(
            degrees=degrees_by_name,
            translation=translation_by_name,
            times=times_by_name,
            image_interpolation=interpolation_by_name,
            **self._get_base_args(),
        )
        transformed = transform(subject)
        assert isinstance(transformed, Subject)
        return transformed

    def get_params(
        self,
        degrees_range: tuple[float, float],
        translation_range: tuple[float, float],
        num_transforms: int,
        perturbation: float = 0.3,
        is_2d: bool = False,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        # If perturbation is 0, time intervals between movements are constant
        degrees_params = self.get_params_array(
            degrees_range,
            num_transforms,
        )
        translation_params = self.get_params_array(
            translation_range,
            num_transforms,
        )
        if is_2d:  # imagine sagittal (1, A, S)
            degrees_params[:, :-1] = 0  # rotate around Z axis only
            translation_params[:, 2] = 0  # translate in XY plane only
        step = 1 / (num_transforms + 1)
        times = torch.arange(0, 1, step)[1:]
        noise = torch.FloatTensor(num_transforms)
        noise.uniform_(-step * perturbation, step * perturbation)
        times += noise
        times_params = times.numpy()
        return times_params, degrees_params, translation_params

    @staticmethod
    def get_params_array(nums_range: tuple[float, float], num_transforms: int):
        tensor = torch.FloatTensor(num_transforms, 3).uniform_(*nums_range)
        return tensor.numpy()


class Motion(IntensityTransform, FourierTransform):
    r"""Add MRI motion artifact.

    Magnetic resonance images suffer from motion artifacts when the subject
    moves during image acquisition. This transform follows
    [Shaw et al., 2019 ](http://proceedings.mlr.press/v102/shaw19a.html) to
    simulate motion artifacts for data augmentation.

    Args:
        degrees: Sequence of rotations $(\theta_1, \theta_2, \theta_3)$.
        translation: Sequence of translations $(t_1, t_2, t_3)$ in mm.
        times: Sequence of times from 0 to 1 at which the motions happen.
        image_interpolation: See Interpolation.
        **kwargs: See [`Transform`][torchio.transforms.Transform] for additional
            keyword arguments.
    """

    def __init__(
        self,
        degrees: np.ndarray | dict[str, np.ndarray],
        translation: np.ndarray | dict[str, np.ndarray],
        times: np.ndarray | dict[str, np.ndarray],
        image_interpolation: str | dict[str, str],
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.degrees = degrees
        self.translation = translation
        self.times = times
        self.image_interpolation = image_interpolation
        self.args_names = [
            'degrees',
            'translation',
            'times',
            'image_interpolation',
        ]

    def apply_transform(self, subject: Subject) -> Subject:
        for image_name, image in self.get_images_dict(subject).items():
            degrees = self.get_parameter(self.degrees, image_name)
            translation = self.get_parameter(self.translation, image_name)
            times = self.get_parameter(self.times, image_name)
            image_interpolation = self.get_parameter(
                self.image_interpolation, image_name
            )
            result_arrays = []
            for channel in image.data:
                sitk_image = nib_to_sitk(
                    channel[np.newaxis],
                    image.affine,
                    force_3d=True,
                )
                transforms = self.get_rigid_transforms(
                    degrees,
                    translation,
                    sitk_image,
                )
                transformed_channel = self.add_artifact(
                    sitk_image,
                    transforms,
                    times,
                    image_interpolation,
                )
                result_arrays.append(transformed_channel)
            result = np.stack(result_arrays)
            image.set_data(torch.as_tensor(result))
        return subject

    def get_rigid_transforms(
        self,
        degrees_params: np.ndarray,
        translation_params: np.ndarray,
        image: sitk.Image,
    ) -> list[sitk.Euler3DTransform]:
        center_ijk = np.array(image.GetSize()) / 2
        center_lps = image.TransformContinuousIndexToPhysicalPoint(center_ijk)
        identity = np.eye(4)
        matrices = [identity]
        zipped = zip(degrees_params, translation_params, strict=True)
        for degrees, translation in zipped:
            radians = np.radians(degrees).tolist()
            motion = sitk.Euler3DTransform()
            motion.SetCenter(center_lps)
            motion.SetRotation(*radians)
            motion.SetTranslation(translation.tolist())
            motion_matrix = self.transform_to_matrix(motion)
            matrices.append(motion_matrix)
        transforms = [self.matrix_to_transform(m) for m in matrices]
        return transforms

    @staticmethod
    def transform_to_matrix(transform: sitk.Euler3DTransform) -> np.ndarray:
        matrix = np.eye(4)
        rotation = np.array(transform.GetMatrix()).reshape(3, 3)
        matrix[:3, :3] = rotation
        matrix[:3, 3] = transform.GetTranslation()
        return matrix

    @staticmethod
    def matrix_to_transform(matrix: np.ndarray) -> sitk.Euler3DTransform:
        transform = sitk.Euler3DTransform()
        rotation = matrix[:3, :3].flatten().tolist()
        transform.SetMatrix(rotation)
        transform.SetTranslation(matrix[:3, 3])
        return transform

    def resample_images(
        self,
        image: sitk.Image,
        transforms: Sequence[sitk.Euler3DTransform],
        interpolation: str,
    ) -> list[sitk.Image]:
        floating = reference = image
        default_value = np.float64(sitk.GetArrayViewFromImage(image).min())
        transforms = transforms[1:]  # first is identity
        images = [image]  # first is identity
        for transform in transforms:
            interpolator = self.get_sitk_interpolator(interpolation)
            resampler = sitk.ResampleImageFilter()
            resampler.SetInterpolator(interpolator)
            resampler.SetReferenceImage(reference)
            resampler.SetOutputPixelType(sitk.sitkFloat32)
            resampler.SetDefaultPixelValue(default_value)
            resampler.SetTransform(transform)
            resampled = resampler.Execute(floating)
            images.append(resampled)
        return images

    @staticmethod
    def sort_spectra(spectra: list[torch.Tensor], times: np.ndarray):
        """Use original spectrum to fill the center of k-space."""
        num_spectra = len(spectra)
        if np.any(times > 0.5):
            index = np.where(times > 0.5)[0].min()
        else:
            index = num_spectra - 1
        spectra[0], spectra[index] = spectra[index], spectra[0]

    def add_artifact(
        self,
        image: sitk.Image,
        transforms: Sequence[sitk.Euler3DTransform],
        times: np.ndarray,
        interpolation: str,
    ):
        images = self.resample_images(image, transforms, interpolation)
        spectra = []
        for image in images:
            array = sitk.GetArrayFromImage(image).transpose()  # sitk to np
            spectrum = self.fourier_transform(torch.from_numpy(array))
            spectra.append(spectrum)
        self.sort_spectra(spectra, times)
        result_spectrum = torch.empty_like(spectra[0])
        last_index = result_spectrum.shape[2]
        indices_array = (last_index * times).astype(int)
        indices = [int(value) for value in indices_array.tolist()]
        indices.append(last_index)
        ini = 0
        for spectrum, fin in zip(spectra, indices, strict=True):
            result_spectrum[..., ini:fin] = spectrum[..., ini:fin]
            ini = fin
        result_image = self.inv_fourier_transform(result_spectrum).real.float()
        return result_image
