#    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 shutil

import nnformer
import numpy as np
from batchgenerators.utilities.file_and_folder_operations import load_pickle, subfiles
from multiprocessing.pool import Pool
from nnformer.configuration import default_num_threads
from nnformer.experiment_planning.common_utils import get_pool_and_conv_props
from nnformer.experiment_planning.experiment_planner_baseline_3DUNet import ExperimentPlanner
from nnformer.experiment_planning.utils import add_classes_in_slice_info
from nnformer.network_architecture.generic_UNet import Generic_UNet
from nnformer.paths import *
from nnformer.preprocessing.preprocessing import PreprocessorFor2D
from nnformer.training.model_restore import recursive_find_python_class


class ExperimentPlanner2D(ExperimentPlanner):
    def __init__(self, folder_with_cropped_data, preprocessed_output_folder):
        super(ExperimentPlanner2D, self).__init__(folder_with_cropped_data,
                                                  preprocessed_output_folder)
        self.data_identifier = default_data_identifier + "_2D"
        self.plans_fname = join(self.preprocessed_output_folder, "nnFormerPlans" + "_plans_2D.pkl")

        self.unet_base_num_features = 30
        self.unet_max_num_filters = 512
        self.unet_max_numpool = 999

        self.preprocessor_name = "PreprocessorFor2D"

    def get_properties_for_stage(self, current_spacing, original_spacing, original_shape, num_cases,
                                 num_modalities, num_classes):

        new_median_shape = np.round(original_spacing / current_spacing * original_shape).astype(int)

        dataset_num_voxels = np.prod(new_median_shape, dtype=np.int64) * num_cases
        input_patch_size = new_median_shape[1:]

        network_numpool, net_pool_kernel_sizes, net_conv_kernel_sizes, input_patch_size, \
        shape_must_be_divisible_by = get_pool_and_conv_props(current_spacing[1:], input_patch_size,
                                                             self.unet_featuremap_min_edge_length,
                                                             self.unet_max_numpool)

        estimated_gpu_ram_consumption = Generic_UNet.compute_approx_vram_consumption(input_patch_size,
                                                                                     network_numpool,
                                                                                     self.unet_base_num_features,
                                                                                     self.unet_max_num_filters,
                                                                                     num_modalities, num_classes,
                                                                                     net_pool_kernel_sizes,
                                                                                     conv_per_stage=self.conv_per_stage)

        batch_size = int(np.floor(Generic_UNet.use_this_for_batch_size_computation_2D /
                                  estimated_gpu_ram_consumption * Generic_UNet.DEFAULT_BATCH_SIZE_2D))
        if batch_size < self.unet_min_batch_size:
            raise RuntimeError("This framework is not made to process patches this large. We will add patch-based "
                               "2D networks later. Sorry for the inconvenience")

        # check if batch size is too large (more than 5 % of dataset)
        max_batch_size = np.round(self.batch_size_covers_max_percent_of_dataset * dataset_num_voxels /
                                  np.prod(input_patch_size, dtype=np.int64)).astype(int)
        batch_size = max(1, min(batch_size, max_batch_size))

        plan = {
            'batch_size': batch_size,
            'num_pool_per_axis': network_numpool,
            'patch_size': input_patch_size,
            'median_patient_size_in_voxels': new_median_shape,
            'current_spacing': current_spacing,
            'original_spacing': original_spacing,
            'pool_op_kernel_sizes': net_pool_kernel_sizes,
            'conv_kernel_sizes': net_conv_kernel_sizes,
            'do_dummy_2D_data_aug': False
        }
        return plan

    def plan_experiment(self):
        if os.path.exists(self.plans_fname):
            self.load_my_plans()
        else:
            use_nonzero_mask_for_normalization = self.determine_whether_to_use_mask_for_norm()
            print("Are we using the nonzero maks for normalizaion?", use_nonzero_mask_for_normalization)

            spacings = self.dataset_properties['all_spacings']
            sizes = self.dataset_properties['all_sizes']
            all_classes = self.dataset_properties['all_classes']
            modalities = self.dataset_properties['modalities']
            num_modalities = len(list(modalities.keys()))

            target_spacing = self.get_target_spacing()
            new_shapes = np.array([np.array(i) / target_spacing * np.array(j) for i, j in zip(spacings, sizes)])

            max_spacing_axis = np.argmax(target_spacing)
            remaining_axes = [i for i in list(range(3)) if i != max_spacing_axis]
            self.transpose_forward = [max_spacing_axis] + remaining_axes
            self.transpose_backward = [np.argwhere(np.array(self.transpose_forward) == i)[0][0] for i in range(3)]

            # we base our calculations on the median shape of the datasets
            median_shape = np.median(np.vstack(new_shapes), 0)
            print("the median shape of the dataset is ", median_shape)

            max_shape = np.max(np.vstack(new_shapes), 0)
            print("the max shape in the dataset is ", max_shape)
            min_shape = np.min(np.vstack(new_shapes), 0)
            print("the min shape in the dataset is ", min_shape)

            print("we don't want feature maps smaller than ", self.unet_featuremap_min_edge_length, " in the bottleneck")

            # how many stages will the image pyramid have?
            self.plans_per_stage = []

            target_spacing_transposed = np.array(target_spacing)[self.transpose_forward]
            median_shape_transposed = np.array(median_shape)[self.transpose_forward]
            print("the transposed median shape of the dataset is ", median_shape_transposed)

            self.plans_per_stage.append(
                self.get_properties_for_stage(target_spacing_transposed, target_spacing_transposed, median_shape_transposed,
                                              num_cases=len(self.list_of_cropped_npz_files),
                                              num_modalities=num_modalities,
                                              num_classes=len(all_classes) + 1),
                )

            print(self.plans_per_stage)

            self.plans_per_stage = self.plans_per_stage[::-1]
            self.plans_per_stage = {i: self.plans_per_stage[i] for i in range(len(self.plans_per_stage))}  # convert to dict

            normalization_schemes = self.determine_normalization_scheme()
            # deprecated
            only_keep_largest_connected_component, min_size_per_class, min_region_size_per_class = None, None, None

            # these are independent of the stage
            plans = {'num_stages': len(list(self.plans_per_stage.keys())), 'num_modalities': num_modalities,
                     'modalities': modalities, 'normalization_schemes': normalization_schemes,
                     'dataset_properties': self.dataset_properties, 'list_of_npz_files': self.list_of_cropped_npz_files,
                     'original_spacings': spacings, 'original_sizes': sizes,
                     'preprocessed_data_folder': self.preprocessed_output_folder, 'num_classes': len(all_classes),
                     'all_classes': all_classes, 'base_num_features': self.unet_base_num_features,
                     'use_mask_for_norm': use_nonzero_mask_for_normalization,
                     'keep_only_largest_region': only_keep_largest_connected_component,
                     'min_region_size_per_class': min_region_size_per_class, 'min_size_per_class': min_size_per_class,
                     'transpose_forward': self.transpose_forward, 'transpose_backward': self.transpose_backward,
                     'data_identifier': self.data_identifier, 'plans_per_stage': self.plans_per_stage,
                     'preprocessor_name': self.preprocessor_name,
                     }

            self.plans = plans
            self.save_my_plans()
