#    Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.
import math
import multiprocessing
import shutil
from time import sleep
from typing import Tuple

import SimpleITK
import numpy as np
import pandas as pd
from batchgenerators.utilities.file_and_folder_operations import *
from tqdm import tqdm
from typing import Union

import nnunetv2
from nnunetv2.paths import nnUNet_preprocessed, nnUNet_raw
from nnunetv2.preprocessing.cropping.cropping import crop_to_nonzero
from nnunetv2.preprocessing.resampling.default_resampling import compute_new_shape
from nnunetv2.training.dataloading.nnunet_dataset import nnUNetDatasetBlosc2
from nnunetv2.utilities.dataset_name_id_conversion import maybe_convert_to_dataset_name
from nnunetv2.utilities.find_class_by_name import recursive_find_python_class
from nnunetv2.utilities.plans_handling.plans_handler import PlansManager, ConfigurationManager
from nnunetv2.utilities.utils import get_filenames_of_train_images_and_targets


class DefaultPreprocessor(object):
    def __init__(self, verbose: bool = True):
        self.verbose = verbose
        """
        Everything we need is in the plans. Those are given when run() is called
        """

    def run_case_npy(self, data: np.ndarray, seg: Union[np.ndarray, None], properties: dict,
                     plans_manager: PlansManager, configuration_manager: ConfigurationManager,
                     dataset_json: Union[dict, str]):
        # let's not mess up the inputs!
        data = data.astype(np.float32)  # this creates a copy
        if seg is not None:
            assert data.shape[1:] == seg.shape[1:], "Shape mismatch between image and segmentation. Please fix your dataset and make use of the --verify_dataset_integrity flag to ensure everything is correct"
            seg = np.copy(seg)

        has_seg = seg is not None

        # apply transpose_forward, this also needs to be applied to the spacing!
        data = data.transpose([0, *[i + 1 for i in plans_manager.transpose_forward]])
        if seg is not None:
            seg = seg.transpose([0, *[i + 1 for i in plans_manager.transpose_forward]])
        original_spacing = [properties['spacing'][i] for i in plans_manager.transpose_forward]

        # crop, remember to store size before cropping!
        shape_before_cropping = data.shape[1:]
        properties['shape_before_cropping'] = shape_before_cropping
        # this command will generate a segmentation. This is important because of the nonzero mask which we may need
        data, seg, bbox = crop_to_nonzero(data, seg)
        properties['bbox_used_for_cropping'] = bbox
        # print(data.shape, seg.shape)
        properties['shape_after_cropping_and_before_resampling'] = data.shape[1:]

        # resample
        target_spacing = configuration_manager.spacing  # this should already be transposed

        if len(target_spacing) < len(data.shape[1:]):
            # target spacing for 2d has 2 entries but the data and original_spacing have three because everything is 3d
            # in 2d configuration we do not change the spacing between slices
            target_spacing = [original_spacing[0]] + target_spacing
        new_shape = compute_new_shape(data.shape[1:], original_spacing, target_spacing)

        # normalize
        # normalization MUST happen before resampling or we get huge problems with resampled nonzero masks no
        # longer fitting the images perfectly!
        data = self._normalize(data, seg, configuration_manager,
                               plans_manager.foreground_intensity_properties_per_channel)

        # print('current shape', data.shape[1:], 'current_spacing', original_spacing,
        #       '\ntarget shape', new_shape, 'target_spacing', target_spacing)
        old_shape = data.shape[1:]
        data = configuration_manager.resampling_fn_data(data, new_shape, original_spacing, target_spacing)
        seg = configuration_manager.resampling_fn_seg(seg, new_shape, original_spacing, target_spacing)
        if self.verbose:
            print(f'old shape: {old_shape}, new_shape: {new_shape}, old_spacing: {original_spacing}, '
                  f'new_spacing: {target_spacing}, fn_data: {configuration_manager.resampling_fn_data}')

        # if we have a segmentation, sample foreground locations for oversampling and add those to properties
        if has_seg:
            # reinstantiating LabelManager for each case is not ideal. We could replace the dataset_json argument
            # with a LabelManager Instance in this function because that's all its used for. Dunno what's better.
            # LabelManager is pretty light computation-wise.
            label_manager = plans_manager.get_label_manager(dataset_json)
            collect_for_this = label_manager.foreground_regions if label_manager.has_regions \
                else label_manager.foreground_labels

            # when using the ignore label we want to sample only from annotated regions. Therefore we also need to
            # collect samples uniformly from all classes (incl background)
            if label_manager.has_ignore_label:
                collect_for_this.append([-1] + label_manager.all_labels)

            # no need to filter background in regions because it is already filtered in handle_labels
            # print(all_labels, regions)
            properties['class_locations'] = self._sample_foreground_locations(seg, collect_for_this,
                                                                                   verbose=self.verbose)
            seg = self.modify_seg_fn(seg, plans_manager, dataset_json, configuration_manager)
        if np.max(seg) > 127:
            seg = seg.astype(np.int16)
        else:
            seg = seg.astype(np.int8)
        return data, seg, properties

    def run_case(self, image_files: List[str], seg_file: Union[str, None], plans_manager: PlansManager,
                 configuration_manager: ConfigurationManager,
                 dataset_json: Union[dict, str]):
        """
        seg file can be none (test cases)

        order of operations is: transpose -> crop -> resample
        so when we export we need to run the following order: resample -> crop -> transpose (we could also run
        transpose at a different place, but reverting the order of operations done during preprocessing seems cleaner)
        """
        if isinstance(dataset_json, str):
            dataset_json = load_json(dataset_json)

        rw = plans_manager.image_reader_writer_class()

        # load image(s)
        data, data_properties = rw.read_images(image_files)

        # if possible, load seg
        if seg_file is not None:
            seg, _ = rw.read_seg(seg_file)
        else:
            seg = None

        if self.verbose:
            print(seg_file)
        data, seg, data_properties = self.run_case_npy(data, seg, data_properties, plans_manager, configuration_manager,
                                      dataset_json)
        return data, seg, data_properties

    def run_case_save(self, output_filename_truncated: str, image_files: List[str], seg_file: str,
                      plans_manager: PlansManager, configuration_manager: ConfigurationManager,
                      dataset_json: Union[dict, str]):
        data, seg, properties = self.run_case(image_files, seg_file, plans_manager, configuration_manager, dataset_json)
        data = data.astype(np.float32, copy=False)
        seg = seg.astype(np.int16, copy=False)
        # print('dtypes', data.dtype, seg.dtype)
        block_size_data, chunk_size_data = nnUNetDatasetBlosc2.comp_blosc2_params(
            data.shape,
            tuple(configuration_manager.patch_size),
            data.itemsize)
        block_size_seg, chunk_size_seg = nnUNetDatasetBlosc2.comp_blosc2_params(
            seg.shape,
            tuple(configuration_manager.patch_size),
            seg.itemsize)

        nnUNetDatasetBlosc2.save_case(data, seg, properties, output_filename_truncated,
                                      chunks=chunk_size_data, blocks=block_size_data,
                                      chunks_seg=chunk_size_seg, blocks_seg=block_size_seg)

    @staticmethod
    def _sample_foreground_locations(
            seg: np.ndarray,
            classes_or_regions: Union[List[int], List[Tuple[int, ...]]],
            seed: int = 1234,
            verbose: bool = False,
            min_num_samples=10000,
            min_percent_coverage = 0.01
    ):

        rndst = np.random.RandomState(seed)

        class_locs = {}

        # Normalize requested labels and compute the set of all labels we might need
        normalized = []
        requested_labels = set()
        for c in classes_or_regions:
            if isinstance(c, (tuple, list)):
                labs = tuple(int(x) for x in c)
                normalized.append(labs)
                requested_labels.update(labs)
            else:
                lab = int(c)
                normalized.append(lab)
                requested_labels.add(lab)

        # Create mask for all requested labels (this includes 0 if requested)
        requested_labels_arr = np.fromiter(requested_labels, dtype=np.int32)
        valid_mask = np.isin(seg, requested_labels_arr)

        coords = np.argwhere(valid_mask)
        seg_sel = seg[valid_mask]
        del valid_mask

        n = seg_sel.size
        if n == 0:
            for c in classes_or_regions:
                k = tuple(c) if isinstance(c, (tuple, list)) else int(c)
                class_locs[k] = []
            return class_locs

        # sort once, then compute label blocks
        order = np.argsort(seg_sel, kind="stable")
        lab_sorted = seg_sel[order]
        coords_sorted = coords[order]

        change = np.flatnonzero(lab_sorted[1:] != lab_sorted[:-1]) + 1
        starts = np.r_[0, change]
        ends = np.r_[change, n]
        labels_present = lab_sorted[starts]

        label_to_range = {int(l): (int(s), int(e)) for l, s, e in zip(labels_present, starts, ends)}
        present_labels = set(label_to_range.keys())

        for c in classes_or_regions:
            is_region = isinstance(c, (tuple, list))
            labs = tuple(int(x) for x in c) if is_region else (int(c),)
            k = labs if is_region else labs[0]

            # Skip if none of the labels are present
            if not any(lab in present_labels for lab in labs):
                class_locs[k] = []
                continue

            # Collect ranges for present labels in this class/region
            ranges = []
            counts = []
            for lab in labs:
                r = label_to_range.get(lab)
                if r is None:
                    continue
                s, e = r
                cnt = e - s
                if cnt > 0:
                    ranges.append((s, e))
                    counts.append(cnt)

            if len(counts) == 0:
                class_locs[k] = []
                continue

            total = int(np.sum(counts))
            target_num_samples = min(min_num_samples, total)
            target_num_samples = max(target_num_samples, int(np.ceil(total * min_percent_coverage)))

            # Sample uniformly without replacement from the union of ranges, without building an n-sized mask
            # Draw target_num_samples unique offsets in [0, total)
            offsets = rndst.choice(total, target_num_samples, replace=False)

            # Map offsets -> (range index, in-range offset) using cumulative counts
            cum = np.cumsum(counts)
            which = np.searchsorted(cum, offsets, side="right")
            prev = np.concatenate(([0], cum[:-1]))
            in_range = offsets - prev[which]

            # Convert to indices in coords_sorted
            starts_for_pick = np.fromiter((ranges[i][0] for i in which), dtype=np.int64, count=which.size)
            picked_idx = starts_for_pick + in_range.astype(np.int64)

            selected = coords_sorted[picked_idx]
            class_locs[k] = selected

            if verbose:
                print(c, target_num_samples)

        return class_locs

    # @staticmethod
    # def _sample_foreground_locations(seg: np.ndarray, classes_or_regions: Union[List[int], List[Tuple[int, ...]]],
    #                                  seed: int = 1234, verbose: bool = False):
    #     num_samples = 10000
    #     min_percent_coverage = 0.01  # at least 1% of the class voxels need to be selected, otherwise it may be too
    #     # sparse
    #     rndst = np.random.RandomState(seed)
    #     class_locs = {}
    #     foreground_mask = seg != 0
    #     foreground_coords = np.argwhere(foreground_mask)
    #     seg = seg[foreground_mask]
    #     del foreground_mask
    #     unique_labels = pd.unique(seg.ravel())
    #
    #     # We don't need more than 1e7 foreground samples. That's insanity. Cap here
    #     if len(foreground_coords) > 1e7:
    #         take_every = math.floor(len(foreground_coords) / 1e7)
    #         # keep computation time reasonable
    #         if verbose:
    #             print(f'Subsampling foreground pixels 1:{take_every} for computational reasons')
    #         foreground_coords = foreground_coords[::take_every]
    #         seg = seg[::take_every]
    #
    #     for c in classes_or_regions:
    #         k = c if not isinstance(c, list) else tuple(c)
    #
    #         # check if any of the labels are in seg, if not skip c
    #         if isinstance(c, (tuple, list)):
    #             if not any([ci in unique_labels for ci in c]):
    #                 class_locs[k] = []
    #                 continue
    #         else:
    #             if c not in unique_labels:
    #                 class_locs[k] = []
    #                 continue
    #
    #         if isinstance(c, (tuple, list)):
    #             mask = seg == c[0]
    #             for cc in c[1:]:
    #                 mask = mask | (seg == cc)
    #             all_locs = foreground_coords[mask]
    #         else:
    #             mask = seg == c
    #             all_locs = foreground_coords[mask]
    #         if len(all_locs) == 0:
    #             class_locs[k] = []
    #             continue
    #         target_num_samples = min(num_samples, len(all_locs))
    #         target_num_samples = max(target_num_samples, int(np.ceil(len(all_locs) * min_percent_coverage)))
    #
    #         selected = all_locs[rndst.choice(len(all_locs), target_num_samples, replace=False)]
    #         class_locs[k] = selected
    #         if verbose:
    #             print(c, target_num_samples)
    #         seg = seg[~mask]
    #         foreground_coords = foreground_coords[~mask]
    #     return class_locs

    def _normalize(self, data: np.ndarray, seg: np.ndarray, configuration_manager: ConfigurationManager,
                   foreground_intensity_properties_per_channel: dict) -> np.ndarray:
        for c in range(data.shape[0]):
            scheme = configuration_manager.normalization_schemes[c]
            normalizer_class = recursive_find_python_class(join(nnunetv2.__path__[0], "preprocessing", "normalization"),
                                                           scheme,
                                                           'nnunetv2.preprocessing.normalization')
            if normalizer_class is None:
                raise RuntimeError(f'Unable to locate class \'{scheme}\' for normalization')
            normalizer = normalizer_class(use_mask_for_norm=configuration_manager.use_mask_for_norm[c],
                                          intensityproperties=foreground_intensity_properties_per_channel[str(c)])
            data[c] = normalizer.run(data[c], seg[0])
        return data

    def run(self, dataset_name_or_id: Union[int, str], configuration_name: str, plans_identifier: str,
            num_processes: int):
        """
        data identifier = configuration name in plans. EZ.
        """
        dataset_name = maybe_convert_to_dataset_name(dataset_name_or_id)

        assert isdir(join(nnUNet_raw, dataset_name)), "The requested dataset could not be found in nnUNet_raw"

        plans_file = join(nnUNet_preprocessed, dataset_name, plans_identifier + '.json')
        assert isfile(plans_file), "Expected plans file (%s) not found. Run corresponding nnUNet_plan_experiment " \
                                   "first." % plans_file
        plans = load_json(plans_file)
        plans_manager = PlansManager(plans)
        configuration_manager = plans_manager.get_configuration(configuration_name)

        if self.verbose:
            print(f'Preprocessing the following configuration: {configuration_name}')
        if self.verbose:
            print(configuration_manager)

        dataset_json_file = join(nnUNet_preprocessed, dataset_name, 'dataset.json')
        dataset_json = load_json(dataset_json_file)

        output_directory = join(nnUNet_preprocessed, dataset_name, configuration_manager.data_identifier)

        if isdir(output_directory):
            shutil.rmtree(output_directory)

        maybe_mkdir_p(output_directory)

        dataset = get_filenames_of_train_images_and_targets(join(nnUNet_raw, dataset_name), dataset_json)

        # identifiers = [os.path.basename(i[:-len(dataset_json['file_ending'])]) for i in seg_fnames]
        # output_filenames_truncated = [join(output_directory, i) for i in identifiers]

        # multiprocessing magic.
        r = []
        with multiprocessing.get_context("spawn").Pool(num_processes) as p:
            remaining = list(range(len(dataset)))
            # p is pretty nifti. If we kill workers they just respawn but don't do any work.
            # So we need to store the original pool of workers.
            workers = [j for j in p._pool]
            for k in dataset.keys():
                r.append(p.starmap_async(self.run_case_save,
                                         ((join(output_directory, k), dataset[k]['images'], dataset[k]['label'],
                                           plans_manager, configuration_manager,
                                           dataset_json),)))

            with tqdm(desc=None, total=len(dataset), disable=self.verbose) as pbar:
                while len(remaining) > 0:
                    all_alive = all([j.is_alive() for j in workers])
                    if not all_alive:
                        raise RuntimeError('Some background worker is 6 feet under. Yuck. \n'
                                           'OK jokes aside.\n'
                                           'One of your background processes is missing. This could be because of '
                                           'an error (look for an error message) or because it was killed '
                                           'by your OS due to running out of RAM. If you don\'t see '
                                           'an error message, out of RAM is likely the problem. In that case '
                                           'reducing the number of workers might help')
                    done = [i for i in remaining if r[i].ready()]
                    # get done so that errors can be raised
                    _ = [r[i].get() for i in done]
                    for _ in done:
                        r[_].get()  # allows triggering errors
                        pbar.update()
                    remaining = [i for i in remaining if i not in done]
                    sleep(0.1)

    def modify_seg_fn(self, seg: np.ndarray, plans_manager: PlansManager, dataset_json: dict,
                      configuration_manager: ConfigurationManager) -> np.ndarray:
        # this function will be called at the end of self.run_case. Can be used to change the segmentation
        # after resampling. Useful for experimenting with sparse annotations: I can introduce sparsity after resampling
        # and don't have to create a new dataset each time I modify my experiments
        return seg


def example_test_case_preprocessing():
    # (paths to files may need adaptations)
    plans_file = '/home/isensee/drives/gpu_data/nnUNet_preprocessed/Dataset219_AMOS2022_postChallenge_task2/nnUNetPlans.json'
    dataset_json_file = '/home/isensee/drives/gpu_data/nnUNet_preprocessed/Dataset219_AMOS2022_postChallenge_task2/dataset.json'
    input_images = ['/home/isensee/drives/e132-rohdaten/nnUNetv2/Dataset219_AMOS2022_postChallenge_task2/imagesTr/amos_0600_0000.nii.gz', ]  # if you only have one channel, you still need a list: ['case000_0000.nii.gz']

    configuration = '3d_fullres'
    pp = DefaultPreprocessor()

    # _ because this position would be the segmentation if seg_file was not None (training case)
    # even if you have the segmentation, don't put the file there! You should always evaluate in the original
    # resolution. What comes out of the preprocessor might have been resampled to some other image resolution (as
    # specified by plans)
    plans_manager = PlansManager(plans_file)
    data, _, properties = pp.run_case(input_images, seg_file=None, plans_manager=plans_manager,
                                      configuration_manager=plans_manager.get_configuration(configuration),
                                      dataset_json=dataset_json_file)

    # voila. Now plug data into your prediction function of choice. We of course recommend nnU-Net's default (TODO)
    return data

def _verify_class_locations(shape, outfile, class_locs):
    import numpy as np
    import SimpleITK as sitk

    out = np.zeros(shape, dtype=np.uint16)  # allow many labels safely

    for i, k in enumerate(class_locs.keys()):
        class_coords = class_locs[k][:, 1:]
        if class_coords is None:
            continue
        class_coords = np.asarray(class_coords)
        if class_coords.size == 0:
            continue

        # Expect coords in (N, 3) as (z, y, x)
        if class_coords.ndim != 2 or class_coords.shape[1] != 3:
            raise ValueError(f"class_locs[{k}] must have shape (N, 3), got {class_coords.shape}")

        z = class_coords[:, 0].astype(np.int64)
        y = class_coords[:, 1].astype(np.int64)
        x = class_coords[:, 2].astype(np.int64)

        # Optional bounds check (cheap and prevents hard-to-debug indexing errors)
        if (z.min() < 0 or y.min() < 0 or x.min() < 0 or
                z.max() >= shape[0] or y.max() >= shape[1] or x.max() >= shape[2]):
            raise ValueError(f"Coordinates for {k} are out of bounds for shape={shape}")

        out[z, y, x] = i + 1  # label 1..K

    img = sitk.GetImageFromArray(out)  # SimpleITK assumes array is z,y,x
    sitk.WriteImage(img, outfile)


if __name__ == '__main__':
    # example_test_case_preprocessing()
    # pp = DefaultPreprocessor()
    # pp.run(2, '2d', 'nnUNetPlans', 8)

    ###########################################################################################################
    # how to process a test cases? This is an example:
    # example_test_case_preprocessing()
    seg = SimpleITK.GetArrayFromImage(SimpleITK.ReadImage('/home/isensee/temp/H-mito-val-v2.nii.gz'))[None]
    a = DefaultPreprocessor._sample_foreground_locations(seg, np.arange(1, np.max(seg) + 1), min_percent_coverage=0.50)

    _verify_class_locations(seg.shape[1:], '/home/isensee/temp/deleteme.nii.gz', a)