from __future__ import annotations

from collections.abc import Generator
from typing import overload

import numpy as np
import torch

from ...constants import MIN_FLOAT_32
from ...types import TypeSpatialShape
from ...types import TypeTripletInt
from ..image import Image
from ..subject import Subject
from .sampler import RandomSampler


class WeightedSampler(RandomSampler):
    r"""Randomly extract patches from a volume given a probability map.

    The probability of sampling a patch centered on a specific voxel is the
    value of that voxel in the probability map. The probabilities need not be
    normalized. For example, voxels can have values 0, 1 and 5. Voxels with
    value 0 will never be at the center of a patch. Voxels with value 5 will
    have 5 times more chance of being at the center of a patch that voxels
    with a value of 1.

    Args:
        patch_size: See [`PatchSampler`](#torchio.data.PatchSampler).
        probability_map: Name of the image in the input subject that will be
            used as a sampling probability map.

    Raises:
        RuntimeError: If the probability map is empty.

    Examples:
        >>> import torchio as tio
        >>> subject = tio.Subject(
        ...     t1=tio.ScalarImage('t1_mri.nii.gz'),
        ...     sampling_map=tio.Image('sampling.nii.gz', type=tio.SAMPLING_MAP),
        ... )
        >>> patch_size = 64
        >>> sampler = tio.data.WeightedSampler(patch_size, 'sampling_map')
        >>> for patch in sampler(subject):
        ...     print(patch[tio.LOCATION])

    Note:
        The index of the center of a patch with even size $s$ is
        arbitrarily set to $s/2$. This is an implementation detail that
        will typically not make any difference in practice.

    Note:
        Values of the probability map near the border will be set to 0 as
        the center of the patch cannot be at the border (unless the patch has
        size 1 or 2 along that axis).
    """

    def __init__(
        self,
        patch_size: TypeSpatialShape,
        probability_map: str | None,
    ):
        super().__init__(patch_size)
        self.probability_map_name = probability_map
        self.cdf = None

    def _generate_patches(
        self,
        subject: Subject,
        num_patches: int | None = None,
    ) -> Generator[Subject]:
        probability_map = self.get_probability_map(subject)
        probability_map_array = self.process_probability_map(
            probability_map,
            subject,
        )
        cdf = self.get_cumulative_distribution_function(probability_map_array)

        patches_left = num_patches if num_patches is not None else True
        while patches_left:
            yield self.extract_patch(subject, probability_map_array, cdf)
            if num_patches is not None:
                patches_left -= 1

    def get_probability_map_image(self, subject: Subject) -> Image:
        assert self.probability_map_name is not None
        if self.probability_map_name in subject:
            return subject.get_image(self.probability_map_name)
        else:
            message = (
                f'Image "{self.probability_map_name}" not found in subject: {subject}'
            )
            raise KeyError(message)

    def get_probability_map(self, subject: Subject) -> torch.Tensor:
        data = self.get_probability_map_image(subject).data
        if torch.any(data < 0):
            message = (
                'Negative values found'
                f' in probability map "{self.probability_map_name}"'
            )
            raise ValueError(message)
        return data

    def process_probability_map(
        self,
        probability_map: torch.Tensor,
        subject: Subject,
    ) -> np.ndarray:
        # Using float32 can create cdf with maximum very far from 1, e.g. 0.92!
        data = probability_map[0].numpy().astype(np.float64)
        assert data.ndim == 3
        self.clear_probability_borders(data, self.patch_size)
        total = data.sum()
        if total == 0:
            half_patch_size = tuple(n // 2 for n in self.patch_size)
            message = (
                'Empty probability map found:'
                f' {self.get_probability_map_image(subject).path}'
                '\nVoxels with positive probability might be near the image'
                ' border.\nIf you suspect that this is the case, try adding a'
                ' padding transform\nwith half the patch size:'
                f' torchio.Pad({half_patch_size})'
            )
            raise RuntimeError(message)
        data /= total  # normalize probabilities
        return data

    @staticmethod
    def clear_probability_borders(
        probability_map: np.ndarray,
        patch_size: np.ndarray,
    ) -> None:
        # Set probability to 0 on voxels that wouldn't possibly be sampled
        # given the current patch size
        # We will arbitrarily define the center of an array with even length
        # using the // Python operator
        # For example, the center of an array (3, 4) will be on (1, 2)
        #
        #   Patch         center
        #  . . . .        . . . .
        #  . . . .   ->   . . x .
        #  . . . .        . . . .
        #
        #
        #    Prob. map      After preprocessing
        #
        #  x x x x x x x       . . . . . . .
        #  x x x x x x x       . . x x x x .
        #  x x x x x x x  -->  . . x x x x .
        #  x x x x x x x  -->  . . x x x x .
        #  x x x x x x x       . . x x x x .
        #  x x x x x x x       . . . . . . .
        #
        # The dots represent removed probabilities, x mark possible locations
        crop_ini = patch_size // 2
        crop_fin = (patch_size - 1) // 2
        crop_i, crop_j, crop_k = crop_ini
        probability_map[:crop_i, :, :] = 0
        probability_map[:, :crop_j, :] = 0
        probability_map[:, :, :crop_k] = 0

        # The call tolist() is very important. Using np.uint16 as negative
        # index will not work because e.g. -np.uint16(2) == 65534
        crop_i, crop_j, crop_k = (int(value) for value in crop_fin.tolist())
        if crop_i:
            probability_map[-crop_i:, :, :] = 0
        if crop_j:
            probability_map[:, -crop_j:, :] = 0
        if crop_k:
            probability_map[:, :, -crop_k:] = 0

    @staticmethod
    def get_cumulative_distribution_function(
        probability_map: np.ndarray,
    ) -> np.ndarray:
        """Return the cumulative distribution function of a probability map."""
        flat_map = probability_map.flatten()
        flat_map_normalized = flat_map / flat_map.sum()
        cdf = np.cumsum(flat_map_normalized)
        return cdf

    @overload
    def extract_patch(
        self,
        subject: Subject,
        index_ini: TypeTripletInt,
        cdf: None = None,
    ) -> Subject: ...

    @overload
    def extract_patch(
        self,
        subject: Subject,
        probability_map: np.ndarray,
        cdf: np.ndarray,
    ) -> Subject: ...

    def extract_patch(
        self,
        subject: Subject,
        index_ini_or_probability_map: TypeTripletInt | np.ndarray,
        cdf: np.ndarray | None = None,
    ) -> Subject:
        if cdf is None:
            if isinstance(index_ini_or_probability_map, np.ndarray):
                index_ini = (
                    int(index_ini_or_probability_map[0]),
                    int(index_ini_or_probability_map[1]),
                    int(index_ini_or_probability_map[2]),
                )
            else:
                index_ini = index_ini_or_probability_map
            return super().extract_patch(subject, index_ini)

        if not isinstance(index_ini_or_probability_map, np.ndarray):
            message = 'Probability map must be a NumPy array when using a CDF'
            raise TypeError(message)

        i, j, k = self.get_random_index_ini(index_ini_or_probability_map, cdf)
        index_ini = i, j, k
        return super().extract_patch(subject, index_ini)

    def get_random_index_ini(
        self,
        probability_map: np.ndarray,
        cdf: np.ndarray,
    ) -> np.ndarray:
        center = self.sample_probability_map(probability_map, cdf)
        assert np.all(center >= 0)
        # See self.clear_probability_borders
        index_ini = center - self.patch_size // 2
        assert np.all(index_ini >= 0)
        return index_ini

    @classmethod
    def sample_probability_map(
        cls,
        probability_map: np.ndarray,
        cdf: np.ndarray,
    ) -> np.ndarray:
        """Inverse transform sampling.

        Examples:
            >>> probability_map = np.array(
            ...    ((0,0,1,1,5,2,1,1,0),
            ...     (2,2,2,2,2,2,2,2,2)))
            >>> probability_map
            array([[0, 0, 1, 1, 5, 2, 1, 1, 0],
                   [2, 2, 2, 2, 2, 2, 2, 2, 2]])
            >>> histogram = np.zeros_like(probability_map)
            >>> for _ in range(100000):
            ...     histogram[WeightedSampler.sample_probability_map(probability_map, cdf)] += 1  # doctest:+SKIP
            ...
            >>> histogram  # doctest:+SKIP
            array([[    0,     0,  3479,  3478, 17121,  7023,  3355,  3378,     0],
                   [ 6808,  6804,  6942,  6809,  6946,  6988,  7002,  6826,  7041]])
        """
        # Get first value larger than random number ensuring the random number
        # is not exactly 0 (see https://github.com/TorchIO-project/torchio/issues/510)
        random_number = max(MIN_FLOAT_32, torch.rand(1).item()) * cdf[-1]

        random_location_index = np.searchsorted(cdf, random_number)

        center = np.unravel_index(
            random_location_index,
            probability_map.shape,
        )

        probability = probability_map[center]
        if probability <= 0:
            message = (
                'Error retrieving probability in weighted sampler.'
                ' Please report this issue at'
                ' https://github.com/TorchIO-project/torchio/issues/new?labels=bug&template=bug_report.md'
            )
            raise RuntimeError(message)

        return np.array(center)
