#!/usr/bin/env python
# Copyright 2018 Division of Medical Image Computing, German Cancer Research Center (DKFZ).
#
# 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.
# ==============================================================================

"""Default Configurations script. Avoids changing configs of all experiments if general settings are to be changed."""

import os

class DefaultConfigs:

    def __init__(self, model, server_env=False, dim=2):
        self.server_env = server_env
        #########################
        #         I/O           #
        #########################

        self.model = model
        self.dim = dim
        # int [0 < dataset_size]. select n patients from dataset for prototyping.
        self.select_prototype_subset = None

        # some default paths.
        self.backbone_path = 'models/backbone.py'
        self.source_dir = os.path.dirname(os.path.realpath(__file__)) #current dir.
        self.input_df_name = 'info_df.pickle'
        self.model_path = 'models/{}.py'.format(self.model)

        if server_env:
            self.source_dir = '/home/jaegerp/code/mamma_code/medicaldetectiontoolkit'

        #########################
        #      Data Loader      #
        #########################

        #random seed for fold_generator and batch_generator.
        self.seed = 0

        #number of threads for multithreaded batch generation.
        self.n_workers = 4 if server_env else os.cpu_count()-1

        # if True, segmentation losses learn all categories, else only foreground vs. background.
        self.class_specific_seg_flag = False

        #########################
        #      Architecture      #
        #########################

        self.weight_decay = 0.0
        # what weight or layer types to exclude from weight decay. options: ["bias", "norm"].
        self.exclude_from_wd = ("norm",)

        # nonlinearity to be applied after convs with nonlinearity. one of 'relu' or 'leaky_relu'
        self.relu = 'relu'

        # if True initializes weights as specified in model script. else use default Pytorch init.
        self.custom_init = False

        # if True adds high-res decoder levels to feature pyramid: P1 + P0. (e.g. set to true in retina_unet configs)
        self.operate_stride1 = False

        #########################
        #  Schedule             #
        #########################

        # number of folds in cross validation.
        self.n_cv_splits = 5


        # number of probabilistic samples in validation.
        self.n_probabilistic_samples = None

        #########################
        #   Testing / Plotting  #
        #########################

        # perform mirroring at test time. (only XY. Z not done to not blow up predictions times).
        self.test_aug = True

        # if True, test data lies in a separate folder and is not part of the cross validation.
        self.hold_out_test_set = True

        # if hold_out_test_set provided, ensemble predictions over models of all trained cv-folds.
        # implications for hold-out test sets: if True, evaluate folds separately on the test set, aggregate only the
        # evaluations. if False, aggregate the raw predictions across all folds, then evaluate.
        self.ensemble_folds = True

        # color specifications for all box_types in prediction_plot.
        self.box_color_palette = {'det': 'b', 'gt': 'r', 'neg_class': 'purple',
                                  'prop': 'w', 'pos_class': 'g', 'pos_anchor': 'c', 'neg_anchor': 'c'}

        # scan over confidence score in evaluation to optimize it on the validation set.
        self.scan_det_thresh = False

        # plots roc-curves / prc-curves in evaluation.
        self.plot_stat_curves = True

        # evaluates average precision per image and averages over images. instead computing one ap over data set.
        self.per_patient_ap = True

        # threshold for clustering 2D box predictions to 3D Cubes. Overlap is computed in XY.
        self.merge_3D_iou = 0.1

        # monitor any value from training.
        self.n_monitoring_figures = 1
        # dict to assign specific plot_values to monitor_figures > 0. {1: ['class_loss'], 2: ['kl_loss', 'kl_sigmas']}
        self.assign_values_to_extra_figure = {}

        # save predictions to csv file in experiment dir.
        self.save_preds_to_csv = True

        # select a maximum number of patient cases to test. number or "all" for all
        self.max_test_patients = "all"



        #########################
        #   MRCNN               #
        #########################

        # if True, mask loss is not applied. used for data sets, where no pixel-wise annotations are provided.
        self.frcnn_mode = True

        # if True, unmolds masks in Mask R-CNN to full-res for plotting/monitoring.
        self.return_masks_in_val = False
        self.return_masks_in_test = False # needed if doing instance segmentation. evaluation not yet implemented.

        # add P6 to Feature Pyramid Network.
        self.sixth_pooling = False

        # for probabilistic detection
        self.n_latent_dims = 0