from __future__ import annotations
from time import time
from typing import Union, List, Tuple, Type

import numpy as np
import torch
from acvl_utils.cropping_and_padding.bounding_boxes import bounding_box_to_slice, insert_crop_into_image
from batchgenerators.utilities.file_and_folder_operations import join

import nnunetv2
from nnunetv2.utilities.find_class_by_name import recursive_find_python_class
from nnunetv2.utilities.helpers import softmax_helper_dim0

from typing import TYPE_CHECKING

# see https://adamj.eu/tech/2021/05/13/python-type-hints-how-to-fix-circular-imports/
if TYPE_CHECKING:
    from nnunetv2.utilities.plans_handling.plans_handler import PlansManager, ConfigurationManager


class LabelManager(object):
    def __init__(self, label_dict: dict, regions_class_order: Union[List[int], None], force_use_labels: bool = False,
                 inference_nonlin=None):
        self._sanity_check(label_dict)
        self.label_dict = label_dict
        self.regions_class_order = regions_class_order
        self._force_use_labels = force_use_labels

        if force_use_labels:
            self._has_regions = False
        else:
            self._has_regions: bool = any(
                [isinstance(i, (tuple, list)) and len(i) > 1 for i in self.label_dict.values()])

        self._ignore_label: Union[None, int] = self._determine_ignore_label()
        self._all_labels: List[int] = self._get_all_labels()

        self._regions: Union[None, List[Union[int, Tuple[int, ...]]]] = self._get_regions()

        if self.has_ignore_label:
            assert self.ignore_label == max(
                self.all_labels) + 1, 'If you use the ignore label it must have the highest ' \
                                      'label value! It cannot be 0 or in between other labels. ' \
                                      'Sorry bro.'

        if inference_nonlin is None:
            self.inference_nonlin = torch.sigmoid if self.has_regions else softmax_helper_dim0
        else:
            self.inference_nonlin = inference_nonlin

    def _sanity_check(self, label_dict: dict):
        if not 'background' in label_dict.keys():
            raise RuntimeError('Background label not declared (remember that this should be label 0!)')
        bg_label = label_dict['background']
        if isinstance(bg_label, (tuple, list)):
            raise RuntimeError(f"Background label must be 0. Not a list. Not a tuple. Your background label: {bg_label}")
        assert int(bg_label) == 0, f"Background label must be 0. Your background label: {bg_label}"
        # not sure if we want to allow regions that contain background. I don't immediately see how this could cause
        # problems so we allow it for now. That doesn't mean that this is explicitly supported. It could be that this
        # just crashes.

    def _get_all_labels(self) -> List[int]:
        all_labels = []
        for k, r in self.label_dict.items():
            # ignore label is not going to be used, hence the name. Duh.
            if k == 'ignore':
                continue
            if isinstance(r, (tuple, list)):
                for ri in r:
                    all_labels.append(int(ri))
            else:
                all_labels.append(int(r))
        all_labels = list(np.unique(all_labels))
        all_labels.sort()
        return all_labels

    def _get_regions(self) -> Union[None, List[Union[int, Tuple[int, ...]]]]:
        if not self._has_regions or self._force_use_labels:
            return None
        else:
            assert self.regions_class_order is not None, 'if region-based training is requested then you need to ' \
                                                         'define regions_class_order!'
            regions = []
            for k, r in self.label_dict.items():
                # ignore ignore label
                if k == 'ignore':
                    continue
                # ignore regions that are background
                if (np.isscalar(r) and r == 0) \
                        or \
                        (isinstance(r, (tuple, list)) and len(np.unique(r)) == 1 and np.unique(r)[0] == 0):
                    continue
                if isinstance(r, list):
                    r = tuple(r)
                regions.append(r)
            assert len(self.regions_class_order) == len(regions), 'regions_class_order must have as ' \
                                                                  'many entries as there are ' \
                                                                  'regions'
            return regions

    def _determine_ignore_label(self) -> Union[None, int]:
        ignore_label = self.label_dict.get('ignore')
        if ignore_label is not None:
            assert isinstance(ignore_label, int), f'Ignore label has to be an integer. It cannot be a region ' \
                                                  f'(list/tuple). Got {type(ignore_label)}.'
        return ignore_label

    @property
    def has_regions(self) -> bool:
        return self._has_regions

    @property
    def has_ignore_label(self) -> bool:
        return self.ignore_label is not None

    @property
    def all_regions(self) -> Union[None, List[Union[int, Tuple[int, ...]]]]:
        return self._regions

    @property
    def all_labels(self) -> List[int]:
        return self._all_labels

    @property
    def ignore_label(self) -> Union[None, int]:
        return self._ignore_label

    def apply_inference_nonlin(self, logits: Union[np.ndarray, torch.Tensor]) -> \
            Union[np.ndarray, torch.Tensor]:
        """
        logits has to have shape (c, x, y(, z)) where c is the number of classes/regions
        """
        if isinstance(logits, np.ndarray):
            logits = torch.from_numpy(logits)

        with torch.no_grad():
            # softmax etc is not implemented for half
            logits = logits.float()
            probabilities = self.inference_nonlin(logits)

        return probabilities

    @torch.inference_mode()
    def convert_probabilities_to_segmentation(self, predicted_probabilities: Union[np.ndarray, torch.Tensor]) -> \
            Union[np.ndarray, torch.Tensor]:
        """
        assumes that inference_nonlinearity was already applied!

        predicted_probabilities has to have shape (c, x, y(, z)) where c is the number of classes/regions
        """
        if not isinstance(predicted_probabilities, (np.ndarray, torch.Tensor)):
            raise RuntimeError(f"Unexpected input type. Expected np.ndarray or torch.Tensor,"
                               f" got {type(predicted_probabilities)}")

        if self.has_regions:
            assert self.regions_class_order is not None, 'if region-based training is requested then you need to ' \
                                                         'define regions_class_order!'
            # check correct number of outputs
        assert predicted_probabilities.shape[0] == self.num_segmentation_heads, \
            f'unexpected number of channels in predicted_probabilities. Expected {self.num_segmentation_heads}, ' \
            f'got {predicted_probabilities.shape[0]}. Remember that predicted_probabilities should have shape ' \
            f'(c, x, y(, z)).'

        if self.has_regions:
            if isinstance(predicted_probabilities, np.ndarray):
                segmentation = np.zeros(predicted_probabilities.shape[1:], dtype=np.uint16)
            else:
                # no uint16 in torch
                segmentation = torch.zeros(predicted_probabilities.shape[1:], dtype=torch.int16,
                                           device=predicted_probabilities.device)
            for i, c in enumerate(self.regions_class_order):
                segmentation[predicted_probabilities[i] > 0.5] = c
        else:
            # numpy is faster than torch. :facepalm:
            is_numpy = isinstance(predicted_probabilities, np.ndarray)
            if not is_numpy:
                predicted_probabilities = predicted_probabilities.numpy()
            segmentation = predicted_probabilities.argmax(0)
            if not is_numpy:
                segmentation = torch.from_numpy(segmentation)

        return segmentation

    @torch.inference_mode()
    def convert_logits_to_segmentation(self, predicted_logits: Union[np.ndarray, torch.Tensor]) -> \
            Union[np.ndarray, torch.Tensor]:
        input_is_numpy = isinstance(predicted_logits, np.ndarray)
        # we can skip this step if we do not have region. Argmax is the same between logits or probabilities
        if self.has_regions:
            probabilities = self.apply_inference_nonlin(predicted_logits)
        else:
            probabilities = predicted_logits
        if input_is_numpy and isinstance(probabilities, torch.Tensor):
            probabilities = probabilities.cpu().numpy()
        return self.convert_probabilities_to_segmentation(probabilities)

    def revert_cropping_on_probabilities(self, predicted_probabilities: Union[torch.Tensor, np.ndarray],
                                         bbox: List[List[int]],
                                         original_shape: Union[List[int], Tuple[int, ...]]):
        """
        ONLY USE THIS WITH PROBABILITIES, DO NOT USE LOGITS AND DO NOT USE FOR SEGMENTATION MAPS!!!

        predicted_probabilities must be (c, x, y(, z))

        Why do we do this here? Well if we pad probabilities we need to make sure that convert_logits_to_segmentation
        correctly returns background in the padded areas. Also we want to ba able to look at the padded probabilities
        and not have strange artifacts.
        Only LabelManager knows how this needs to be done. So let's let him/her do it, ok?
        """
        # revert cropping
        probs_reverted_cropping = np.zeros((predicted_probabilities.shape[0], *original_shape),
                                           dtype=predicted_probabilities.dtype) \
            if isinstance(predicted_probabilities, np.ndarray) else \
            torch.zeros((predicted_probabilities.shape[0], *original_shape), dtype=predicted_probabilities.dtype)

        if not self.has_regions:
            probs_reverted_cropping[0] = 1

        probs_reverted_cropping = insert_crop_into_image(probs_reverted_cropping, predicted_probabilities, bbox)
        return probs_reverted_cropping

    @staticmethod
    def filter_background(classes_or_regions: Union[List[int], List[Union[int, Tuple[int, ...]]]]):
        # heck yeah
        # This is definitely taking list comprehension too far. Enjoy.
        return [i for i in classes_or_regions if
                ((not isinstance(i, (tuple, list))) and i != 0)
                or
                (isinstance(i, (tuple, list)) and not (
                        len(np.unique(i)) == 1 and np.unique(i)[0] == 0))]

    @property
    def foreground_regions(self):
        return self.filter_background(self.all_regions)

    @property
    def foreground_labels(self):
        return self.filter_background(self.all_labels)

    @property
    def num_segmentation_heads(self):
        if self.has_regions:
            return len(self.foreground_regions)
        else:
            return len(self.all_labels)


def get_labelmanager_class_from_plans(plans: dict) -> Type[LabelManager]:
    if 'label_manager' not in plans.keys():
        print('No label manager specified in plans. Using default: LabelManager')
        return LabelManager
    else:
        labelmanager_class = recursive_find_python_class(join(nnunetv2.__path__[0], "utilities", "label_handling"),
                                                         plans['label_manager'],
                                                         current_module="nnunetv2.utilities.label_handling")
        return labelmanager_class


def convert_labelmap_to_one_hot(segmentation: Union[np.ndarray, torch.Tensor],
                                all_labels: Union[List, torch.Tensor, np.ndarray, tuple],
                                output_dtype=None) -> Union[np.ndarray, torch.Tensor]:
    """
    if output_dtype is None then we use np.uint8/torch.uint8
    if input is torch.Tensor then output will be on the same device

    np.ndarray is faster than torch.Tensor

    if segmentation is torch.Tensor, this function will be faster if it is LongTensor. If it is somethine else we have
    to cast which takes time.

    IMPORTANT: This function only works properly if your labels are consecutive integers, so something like 0, 1, 2, 3, ...
    DO NOT use it with 0, 32, 123, 255, ... or whatever (fix your labels, yo)
    """
    if isinstance(segmentation, torch.Tensor):
        result = torch.zeros((len(all_labels), *segmentation.shape),
                             dtype=output_dtype if output_dtype is not None else (torch.uint8 if max(all_labels) < 255 else torch.uint16),
                             device=segmentation.device)
        # variant 1, 2x faster than 2
        result.scatter_(0, segmentation[None].long(), 1)  # why does this have to be long!?
        # variant 2, slower than 1
        # for i, l in enumerate(all_labels):
        #     result[i] = segmentation == l
    else:
        result = np.zeros((len(all_labels), *segmentation.shape),
                          dtype=output_dtype if output_dtype is not None else (np.uint8 if max(all_labels) < 255 else np.uint16))
        # variant 1, fastest in my testing
        for i, l in enumerate(all_labels):
            result[i] = segmentation == l
        # variant 2. Takes about twice as long so nah
        # result = np.eye(len(all_labels))[segmentation].transpose((3, 0, 1, 2))
    return result


def determine_num_input_channels(plans_manager: PlansManager,
                                 configuration_or_config_manager: Union[str, ConfigurationManager],
                                 dataset_json: dict) -> int:
    if isinstance(configuration_or_config_manager, str):
        config_manager = plans_manager.get_configuration(configuration_or_config_manager)
    else:
        config_manager = configuration_or_config_manager

    label_manager = plans_manager.get_label_manager(dataset_json)
    num_modalities = len(dataset_json['modality']) if 'modality' in dataset_json.keys() else len(dataset_json['channel_names'])

    # cascade has different number of input channels
    if config_manager.previous_stage_name is not None:
        num_label_inputs = len(label_manager.foreground_labels)
        num_input_channels = num_modalities + num_label_inputs
    else:
        num_input_channels = num_modalities
    return num_input_channels


if __name__ == '__main__':
    # this code used to be able to differentiate variant 1 and 2 to measure time.
    num_labels = 7
    seg = np.random.randint(0, num_labels, size=(256, 256, 256), dtype=np.uint8)
    seg_torch = torch.from_numpy(seg)
    st = time()
    onehot_npy = convert_labelmap_to_one_hot(seg, np.arange(num_labels))
    time_1 = time()
    onehot_npy2 = convert_labelmap_to_one_hot(seg, np.arange(num_labels))
    time_2 = time()
    onehot_torch = convert_labelmap_to_one_hot(seg_torch, np.arange(num_labels))
    time_torch = time()
    onehot_torch2 = convert_labelmap_to_one_hot(seg_torch, np.arange(num_labels))
    time_torch2 = time()
    print(
        f'np: {time_1 - st}, np2: {time_2 - time_1}, torch: {time_torch - time_2}, torch2: {time_torch2 - time_torch}')
    onehot_torch = onehot_torch.numpy()
    onehot_torch2 = onehot_torch2.numpy()
    print(np.all(onehot_torch == onehot_npy))
    print(np.all(onehot_torch2 == onehot_npy))
