from __future__ import annotations

from collections.abc import Sequence

import numpy as np
import torch

from ....data.subject import Subject
from ....types import TypeTripletInt
from ....types import TypeTuple
from ....utils import to_tuple
from ...intensity_transform import IntensityTransform
from .. import RandomTransform

TypeLocations = Sequence[tuple[TypeTripletInt, TypeTripletInt]]


class RandomSwap(RandomTransform, IntensityTransform):
    r"""Randomly swap patches within an image.

    This is typically used in [context restoration for self-supervised learning
    ](https://www.sciencedirect.com/science/article/pii/S1361841518304699).

    Args:
        patch_size: Tuple of integers $(w, h, d)$ to swap patches
            of size $w \times h \times d$.
            If a single number $n$ is provided, $w = h = d = n$.
        num_iterations: Number of times that two patches will be swapped.
        **kwargs: See [`Transform`][torchio.transforms.Transform] for additional
            keyword arguments.
    """

    def __init__(
        self,
        patch_size: TypeTuple = 15,
        num_iterations: int = 100,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.patch_size = np.array(to_tuple(patch_size))
        self.num_iterations = self._parse_num_iterations(num_iterations)

    @staticmethod
    def _parse_num_iterations(num_iterations):
        if not isinstance(num_iterations, int):
            raise TypeError(
                f'num_iterations must be an int,not {num_iterations}',
            )
        if num_iterations < 0:
            raise ValueError(
                f'num_iterations must be positive,not {num_iterations}',
            )
        return num_iterations

    @staticmethod
    def get_params(
        tensor: torch.Tensor,
        patch_size: np.ndarray,
        num_iterations: int,
    ) -> list[tuple[TypeTripletInt, TypeTripletInt]]:
        si, sj, sk = tensor.shape[-3:]
        spatial_shape = si, sj, sk  # for mypy
        patch_size_values = [int(value) for value in patch_size.tolist()]
        locations = []
        for _ in range(num_iterations):
            first_ini, first_fin = get_random_indices_from_shape(
                spatial_shape,
                patch_size_values,
            )
            while True:
                second_ini, second_fin = get_random_indices_from_shape(
                    spatial_shape,
                    patch_size_values,
                )
                larger_than_initial = np.all(second_ini >= first_ini)
                less_than_final = np.all(second_fin <= first_fin)
                if larger_than_initial and less_than_final:
                    continue  # patches overlap
                else:
                    break  # patches don't overlap
            first_values = [int(value) for value in first_ini.tolist()]
            second_values = [int(value) for value in second_ini.tolist()]
            first_location: TypeTripletInt = (
                first_values[0],
                first_values[1],
                first_values[2],
            )
            second_location: TypeTripletInt = (
                second_values[0],
                second_values[1],
                second_values[2],
            )
            location = first_location, second_location
            locations.append(location)
        return locations

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

        locations_by_name: dict[str, TypeLocations] = {}
        patch_size_by_name: dict[str, TypeTripletInt] = {}
        broadcast = np.broadcast_to(self.patch_size, 3)
        patch_size = (int(broadcast[0]), int(broadcast[1]), int(broadcast[2]))
        for name, image in images_dict.items():
            locations = self.get_params(
                image.data,
                self.patch_size,
                self.num_iterations,
            )
            locations_by_name[name] = locations
            patch_size_by_name[name] = patch_size
        transform = Swap(
            patch_size=patch_size_by_name,
            locations=locations_by_name,
            **self._get_base_args(),
        )
        transformed = transform(subject)
        assert isinstance(transformed, Subject)
        return transformed


class Swap(IntensityTransform):
    r"""Swap patches within an image.

    This is typically used in [context restoration for self-supervised learning
    ](https://www.sciencedirect.com/science/article/pii/S1361841518304699).

    Args:
        patch_size: Tuple of integers $(w, h, d)$ to swap patches
            of size $w \times h \times d$.
            If a single number $n$ is provided, $w = h = d = n$.
        num_iterations: Number of times that two patches will be swapped.
        **kwargs: See [`Transform`][torchio.transforms.Transform] for additional
            keyword arguments.
    """

    def __init__(
        self,
        patch_size: TypeTripletInt | dict[str, TypeTripletInt],
        locations: TypeLocations | dict[str, TypeLocations],
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.locations = locations
        self.patch_size = patch_size
        self.args_names = ['locations', 'patch_size']
        self.invert_transform = False

    def apply_transform(self, subject: Subject) -> Subject:
        for name, image in self.get_images_dict(subject).items():
            locations = list(self.get_parameter(self.locations, name))
            patch_size = self.get_parameter(self.patch_size, name)
            if self.invert_transform:
                locations.reverse()
            swapped = _swap(image.data, patch_size, locations)
            image.set_data(swapped)
        return subject


def _swap(
    tensor: torch.Tensor,
    patch_size: TypeTuple,
    locations: TypeLocations,
) -> torch.Tensor:
    # Note this function modifies the input in-place
    tensor = tensor.clone()
    patch_size_array = np.array(patch_size)
    for first_ini, second_ini in locations:
        first_ini_array = np.asarray(first_ini)
        second_ini_array = np.asarray(second_ini)
        first_fin = first_ini_array + patch_size_array
        second_fin = second_ini_array + patch_size_array
        first_patch = _crop(tensor, first_ini_array, first_fin)
        second_patch = _crop(tensor, second_ini_array, second_fin).clone()
        _insert(tensor, first_patch, second_ini_array)
        _insert(tensor, second_patch, first_ini_array)
    return tensor


def _insert(
    tensor: torch.Tensor,
    patch: torch.Tensor,
    index_ini: np.ndarray,
) -> None:
    index_fin = index_ini + np.array(patch.shape[-3:])
    i_ini, j_ini, k_ini = index_ini
    i_fin, j_fin, k_fin = index_fin
    tensor[:, i_ini:i_fin, j_ini:j_fin, k_ini:k_fin] = patch


def _crop(
    image: torch.Tensor,
    index_ini: np.ndarray,
    index_fin: np.ndarray,
) -> torch.Tensor:
    i_ini, j_ini, k_ini = index_ini
    i_fin, j_fin, k_fin = index_fin
    return image[:, i_ini:i_fin, j_ini:j_fin, k_ini:k_fin]


def get_random_indices_from_shape(
    spatial_shape: Sequence[int],
    patch_size: Sequence[int],
) -> tuple[np.ndarray, np.ndarray]:
    assert len(spatial_shape) == 3
    assert len(patch_size) in (1, 3)
    shape_array = np.array(spatial_shape)
    patch_size_array = np.array(patch_size)
    max_index_ini_unchecked = shape_array - patch_size_array
    if (max_index_ini_unchecked < 0).any():
        message = (
            f'Patch size {patch_size} cannot be'
            f' larger than image spatial shape {spatial_shape}'
        )
        raise ValueError(message)
    max_index_ini = max_index_ini_unchecked.astype(np.uint16)
    coordinates = []
    for max_coordinate in max_index_ini.tolist():
        if max_coordinate == 0:
            coordinate = 0
        else:
            coordinate = int(torch.randint(max_coordinate, size=(1,)).item())
        coordinates.append(coordinate)
    index_ini = np.array(coordinates, np.uint16)
    index_fin = index_ini + patch_size_array
    return index_ini, index_fin
