#!/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.

# ==============================================================================

import os

import code

import numpy as np

import torch

from scipy.stats import norm

from collections import OrderedDict

from multiprocessing import Pool

import pickle

from copy import deepcopy

import pandas as pd

import utils.exp_utils as utils

from plotting import plot_batch_prediction

class Predictor:

    """

    Prediction pipeline:

    - receives a patched patient image (n_patches, c, y, x, (z)) from patient data loader.

    - forwards patches through model in chunks of batch_size. (method: batch_tiling_forward)

    - unmolds predictions (boxes and segmentations) to original patient coordinates. (method: spatial_tiling_forward)

    Ensembling (mode == 'test'):

    - for inference, forwards 4 mirrored versions of image to through model and unmolds predictions afterwards

      accordingly (method: data_aug_forward)

    - for inference, loads multiple parameter-sets of the trained model corresponding to different epochs. for each

      parameter-set loops over entire test set, runs prediction pipeline for each patient. (method: predict_test_set)

    Consolidation of predictions:

    - consolidates a patient's predictions (boxes, segmentations) collected over patches, data_aug- and temporal ensembling,

      performs clustering and weighted averaging (external function: apply_wbc_to_patient) to obtain consistent outptus.

    - for 2D networks, consolidates box predictions to 3D cubes via clustering (adaption of non-maximum surpression).

      (external function: merge_2D_to_3D_preds_per_patient)

    Ground truth handling:

    - dissmisses any ground truth boxes returned by the model (happens in validation mode, patch-based groundtruth)

    - if provided by data loader, adds 3D ground truth to the final predictions to be passed to the evaluator.

    """

    def __init__(self, cf, net, logger, mode):

        self.cf = cf

        self.logger = logger

        # mode is 'val' for patient-based validation/monitoring and 'test' for inference.

        self.mode = mode

        # model instance. In validation mode, contains parameters of current epoch.

        self.net = net

        # rank of current epoch loaded (for temporal averaging). this info is added to each prediction,

        # for correct weighting during consolidation.

        self.rank_ix = '0'

        # number of ensembled models. used to calculate the number of expected predictions per position

        # during consolidation of predictions. Default is 1 (no ensembling, e.g. in validation).

        self.n_ens = 1

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        if self.mode == 'test':

            try:

                self.epoch_ranking = np.load(os.path.join(self.cf.fold_dir, 'epoch_ranking.npy'))[:cf.test_n_epochs]

            except:

                raise RuntimeError('no epoch ranking file in fold directory. '

                                   'seems like you are trying to run testing without prior training...')

            self.n_ens = cf.test_n_epochs

            if self.cf.test_aug:

                self.n_ens *= 4

            self.example_plot_dir = os.path.join(cf.test_dir, "example_plots")

            os.makedirs(self.example_plot_dir, exist_ok=True)

    def predict_patient(self, batch):

        """

        predicts one patient.

        called either directly via loop over validation set in exec.py (mode=='val')

        or from self.predict_test_set (mode=='test).

        in val mode:  adds 3D ground truth info to predictions and runs consolidation and 2Dto3D merging of predictions.

        in test mode: returns raw predictions (ground truth addition, consolidation, 2D to 3D merging are

                      done in self.predict_test_set, because patient predictions across several epochs might be needed

                      to be collected first, in case of temporal ensembling).

        :return. results_dict: stores the results for one patient. dictionary with keys:

                 - 'boxes': list over batch elements. each element is a list over boxes, where each box is

                            one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions

                            (if not merged to 3D), and a dummy batch dimension of 1 for 3D predictions.

                 - 'seg_preds': pixel-wise predictions. (b, 1, y, x, (z))

                 - losses (only in validation mode)

        """

        #self.logger.info('\revaluating patient {} for fold {} '.format(batch['pid'], self.cf.fold))

        print('\revaluating patient {} for fold {} '.format(batch['pid'], self.cf.fold), end="", flush=True)

        # True if patient is provided in patches and predictions need to be tiled.

        self.patched_patient = 'patch_crop_coords' in batch.keys()

        # forward batch through prediction pipeline.

        results_dict = self.data_aug_forward(batch)

        if self.mode == 'val':

            for b in range(batch['patient_bb_target'].shape[0]):

                for t in range(len(batch['patient_bb_target'][b])):

                    results_dict['boxes'][b].append({'box_coords': batch['patient_bb_target'][b][t],

                                                     'box_label': batch['patient_roi_labels'][b][t],

                                                     'box_type': 'gt'})

            if self.patched_patient:

                wcs_input = [results_dict['boxes'], 'dummy_pid', self.cf.class_dict, self.cf.wcs_iou, self.n_ens]

                results_dict['boxes'] = apply_wbc_to_patient(wcs_input)[0]

            if self.cf.merge_2D_to_3D_preds:

                merge_dims_inputs = [results_dict['boxes'], 'dummy_pid', self.cf.class_dict, self.cf.merge_3D_iou]

                results_dict['boxes'] = merge_2D_to_3D_preds_per_patient(merge_dims_inputs)[0]

        return results_dict

    def predict_test_set(self, batch_gen, return_results=True):

        """

        wrapper around test method, which loads multiple (or one) epoch parameters (temporal ensembling), loops through

        the test set and collects predictions per patient. Also flattens the results per patient and epoch

        and adds optional ground truth boxes for evaluation. Saves out the raw result list for later analysis and

        optionally consolidates and returns predictions immediately.

        :return: (optionally) list_of_results_per_patient: list over patient results. each entry is a dict with keys:

                 - 'boxes': list over batch elements. each element is a list over boxes, where each box is

                            one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions

                            (if not merged to 3D), and a dummy batch dimension of 1 for 3D predictions.

                 - 'seg_preds': not implemented yet. todo for evaluation of instance/semantic segmentation.

        """

        dict_of_patient_results = OrderedDict()

        # get paths of all parameter sets to be loaded for temporal ensembling. (or just one for no temp. ensembling).

        weight_paths = [os.path.join(self.cf.fold_dir, '{}_best_checkpoint'.format(epoch), 'params.pth') for epoch in

                        self.epoch_ranking]

        print (weight_paths)

        n_test_plots = min(batch_gen['n_test'], 1)

        for rank_ix, weight_path in enumerate(weight_paths):

            self.logger.info(('tmp ensembling over rank_ix:{} epoch:{}'.format(rank_ix, weight_path)))

            #code.interact(local=locals())

            self.net.load_state_dict(torch.load(weight_path))

            self.net.eval()

            self.rank_ix = str(rank_ix)  # get string of current rank for unique patch ids.

            plot_batches = np.random.choice(np.arange(batch_gen['n_test']), size=n_test_plots, replace=False)

            with torch.no_grad():

                for i in range(batch_gen['n_test']):

                    batch = next(batch_gen['test'])

                    # store batch info in patient entry of results dict.

                    if rank_ix == 0:

                        dict_of_patient_results[batch['pid']] = {}

                        dict_of_patient_results[batch['pid']]['results_dicts'] = []

                        dict_of_patient_results[batch['pid']]['patient_bb_target'] = batch['patient_bb_target']

                        dict_of_patient_results[batch['pid']]['patient_roi_labels'] = batch['patient_roi_labels']

                    # call prediction pipeline and store results in dict.

                    results_dict = self.predict_patient(batch)

                    dict_of_patient_results[batch['pid']]['results_dicts'].append({"boxes": results_dict['boxes']})

                    if i in plot_batches and not self.patched_patient:

                        # view qualitative results of random test case

                        # plotting for patched patients is too expensive, thus not done. Change at will.

                        try:

                            out_file = os.path.join(self.example_plot_dir,

                                                    'batch_example_test_{}_rank_{}.png'.format(self.cf.fold,

                                                                                               rank_ix))

                            results_for_plotting = deepcopy(results_dict)

                            # seg preds of test augs are included separately. for viewing, only show aug 0 (merging

                            # would need multiple changes, incl in every model).

                            if results_for_plotting["seg_preds"].shape[1] > 1:

                                results_for_plotting["seg_preds"] = results_dict['seg_preds'][:, [0]]

                            for bix in range(batch["seg"].shape[0]): # batch dim should be 1

                                for tix in range(len(batch['bb_target'][bix])):

                                    results_for_plotting['boxes'][bix].append({'box_coords': batch['bb_target'][bix][tix],

                                                                       'box_label': batch['class_target'][bix][tix],

                                                                       'box_type': 'gt'})

                            utils.split_off_process(plot_batch_prediction, batch, results_for_plotting, self.cf,

                                                    outfile=out_file, suptitle="Test plot:\nunmerged TTA overlayed.")

                        except Exception as e:

                            self.logger.info("WARNING: error in plotting example test batch: {}".format(e))

        self.logger.info('finished predicting test set. starting post-processing of predictions.')

        results_per_patient = []

        # loop over patients again to flatten results across epoch predictions.

        # if provided, add ground truth boxes for evaluation.

        for pid, p_dict in dict_of_patient_results.items():

            tmp_ens_list = p_dict['results_dicts']

            results_dict = {}

            # collect all boxes/seg_preds of same batch_instance over temporal instances.

            b_size = len(tmp_ens_list[0]["boxes"])

            results_dict['boxes'] = [[item for rank_dict in tmp_ens_list for item in rank_dict["boxes"][batch_instance]]

                                     for batch_instance in range(b_size)]

            # TODO return for instance segmentation:

            # results_dict['seg_preds'] = np.mean(results_dict['seg_preds'], 1)[:, None]

            # results_dict['seg_preds'] = np.array([[item for d in tmp_ens_list for item in d['seg_preds'][batch_instance]]

            #                                       for batch_instance in range(len(tmp_ens_list[0]['boxes']))])

            # add 3D ground truth boxes for evaluation.

            for b in range(p_dict['patient_bb_target'].shape[0]):

                for t in range(len(p_dict['patient_bb_target'][b])):

                    results_dict['boxes'][b].append({'box_coords': p_dict['patient_bb_target'][b][t],

                                                     'box_label': p_dict['patient_roi_labels'][b][t],

                                                     'box_type': 'gt'})

            results_per_patient.append([results_dict, pid])

        # save out raw predictions.

        out_string = 'raw_pred_boxes_hold_out_list' if self.cf.hold_out_test_set else 'raw_pred_boxes_list'

        with open(os.path.join(self.cf.fold_dir, '{}.pickle'.format(out_string)), 'wb') as handle:

            pickle.dump(results_per_patient, handle)

        if return_results:

            final_patient_box_results = [(res_dict["boxes"], pid) for res_dict, pid in results_per_patient]

            # consolidate predictions.

            self.logger.info('applying wcs to test set predictions with iou = {} and n_ens = {}.'.format(

                self.cf.wcs_iou, self.n_ens))

            pool = Pool(processes=8)

            mp_inputs = [[ii[0], ii[1], self.cf.class_dict, self.cf.wcs_iou, self.n_ens] for ii in final_patient_box_results]

            final_patient_box_results = pool.map(apply_wbc_to_patient, mp_inputs, chunksize=1)

            pool.close()

            pool.join()

            # merge 2D boxes to 3D cubes. (if model predicts 2D but evaluation is run in 3D)

            if self.cf.merge_2D_to_3D_preds:

                self.logger.info('applying 2Dto3D merging to test set predictions with iou = {}.'.format(self.cf.merge_3D_iou))

                pool = Pool(processes=6)

                mp_inputs = [[ii[0], ii[1], self.cf.class_dict, self.cf.merge_3D_iou] for ii in final_patient_box_results]

                final_patient_box_results = pool.map(merge_2D_to_3D_preds_per_patient, mp_inputs, chunksize=1)

                pool.close()

                pool.join()

            # final_patient_box_results holds [avg_boxes, pid] if wbc

            for ix in range(len(results_per_patient)):

                assert results_per_patient[ix][1] == final_patient_box_results[ix][1], "should be same pid"

                results_per_patient[ix][0]["boxes"] = final_patient_box_results[ix][0]

            out_string = 'wbc_pred_boxes_hold_out_list' if self.cf.hold_out_test_set else 'wbc_pred_boxes_list'

            with open(os.path.join(self.cf.fold_dir, '{}.pickle'.format(out_string)), 'wb') as handle:

               pickle.dump(results_per_patient, handle)

            return results_per_patient

    def load_saved_predictions(self, apply_wbc=False):

        """

        loads raw predictions saved by self.predict_test_set. consolidates and merges 2D boxes to 3D cubes for evaluation.

        (if model predicts 2D but evaluation is run in 3D)

        :return: (optionally) results_list: list over patient results. each entry is a dict with keys:

                 - 'boxes': list over batch elements. each element is a list over boxes, where each box is

                            one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions

                            (if not merged to 3D), and a dummy batch dimension of 1 for 3D predictions.

                 - 'seg_preds': not implemented yet. todo for evaluation of instance/semantic segmentation.

        """

        # load predictions for a single test-set fold.

        results_file = 'raw_pred_boxes_hold_out_list.pickle' if self.cf.hold_out_test_set else 'raw_pred_boxes_list.pickle'

        if not self.cf.hold_out_test_set or not self.cf.ensemble_folds:

            with open(os.path.join(self.cf.fold_dir, results_file), 'rb') as handle:

                results_list = pickle.load(handle)

            box_results_list = [(res_dict["boxes"], pid) for res_dict, pid in results_list]

            da_factor = 4 if self.cf.test_aug else 1

            n_ens = self.cf.test_n_epochs * da_factor

            self.logger.info('loaded raw test set predictions with n_patients = {} and n_ens = {}'.format(

                len(results_list), n_ens))

        # if hold out test set was perdicted, aggregate predictions of all trained models

        # corresponding to all CV-folds and flatten them.

        else:

            self.logger.info("loading saved predictions of hold-out test set and ensembling over folds.")

            fold_dirs = sorted([os.path.join(self.cf.exp_dir, f) for f in os.listdir(self.cf.exp_dir) if

                                os.path.isdir(os.path.join(self.cf.exp_dir, f)) and f.startswith("fold")])

            results_list = []

            folds_loaded = 0

            for fold in range(self.cf.n_cv_splits):

                fold_dir = os.path.join(self.cf.exp_dir, 'fold_{}'.format(fold))

                if fold_dir in fold_dirs:

                    with open(os.path.join(fold_dir, results_file), 'rb') as handle:

                        fold_list = pickle.load(handle)

                        results_list += fold_list

                        folds_loaded += 1

                else:

                    self.logger.info("Skipping fold {} since no saved predictions found.".format(fold))

            box_results_list = []

            for res_dict, pid in results_list: #without filtering gt out:

                box_results_list.append((res_dict['boxes'], pid))

            da_factor = 4 if self.cf.test_aug else 1

            n_ens = self.cf.test_n_epochs * da_factor * folds_loaded

        # consolidate predictions.

        if apply_wbc:

            self.logger.info('applying wcs to test set predictions with iou = {} and n_ens = {}.'.format(

                self.cf.wcs_iou, n_ens))

            pool = Pool(processes=6)

            mp_inputs = [[ii[0], ii[1], self.cf.class_dict, self.cf.wcs_iou, n_ens] for ii in box_results_list]

            box_results_list = pool.map(apply_wbc_to_patient, mp_inputs, chunksize=1)

            pool.close()

            pool.join()

        # merge 2D box predictions to 3D cubes (if model predicts 2D but evaluation is run in 3D)

        if self.cf.merge_2D_to_3D_preds:

            self.logger.info(

                'applying 2Dto3D merging to test set predictions with iou = {}.'.format(self.cf.merge_3D_iou))

            pool = Pool(processes=6)

            mp_inputs = [[ii[0], ii[1], self.cf.class_dict, self.cf.merge_3D_iou] for ii in box_results_list]

            box_results_list = pool.map(merge_2D_to_3D_preds_per_patient, mp_inputs, chunksize=1)

            pool.close()

            pool.join()

        for ix in range(len(results_list)):

            assert np.all(

                results_list[ix][1] == box_results_list[ix][1]), "pid mismatch between loaded and aggregated results"

            results_list[ix][0]["boxes"] = box_results_list[ix][0]

        return results_list  # holds (results_dict, pid)

    def data_aug_forward(self, batch):

        """

        in val_mode: passes batch through to spatial_tiling method without data_aug.

        in test_mode: if cf.test_aug is set in configs, createst 4 mirrored versions of the input image,

        passes all of them to the next processing step (spatial_tiling method) and re-transforms returned predictions

        to original image version.

        :return. results_dict: stores the results for one patient. dictionary with keys:

                 - 'boxes': list over batch elements. each element is a list over boxes, where each box is

                            one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions,

                            and a dummy batch dimension of 1 for 3D predictions.

                 - 'seg_preds': pixel-wise predictions. (b, 1, y, x, (z))

                 - losses (only in validation mode)

        """

        patch_crops = batch['patch_crop_coords'] if self.patched_patient else None

        results_list = [self.spatial_tiling_forward(batch, patch_crops)]

        org_img_shape = batch['original_img_shape']

        if self.mode == 'test' and self.cf.test_aug:

            if self.patched_patient:

                # apply mirror transformations to patch-crop coordinates, for correct tiling in spatial_tiling method.

                mirrored_patch_crops = get_mirrored_patch_crops(patch_crops, batch['original_img_shape'])

            else:

                mirrored_patch_crops = [None] * 3

            img = np.copy(batch['data'])

            # first mirroring: y-axis.

            batch['data'] = np.flip(img, axis=2).copy()

            chunk_dict = self.spatial_tiling_forward(batch, mirrored_patch_crops[0], n_aug='1')

            # re-transform coordinates.

            for ix in range(len(chunk_dict['boxes'])):

                for boxix in range(len(chunk_dict['boxes'][ix])):

                    coords = chunk_dict['boxes'][ix][boxix]['box_coords'].copy()

                    coords[0] = org_img_shape[2] - chunk_dict['boxes'][ix][boxix]['box_coords'][2]

                    coords[2] = org_img_shape[2] - chunk_dict['boxes'][ix][boxix]['box_coords'][0]

                    assert coords[2] >= coords[0], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]

                    assert coords[3] >= coords[1], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]

                    chunk_dict['boxes'][ix][boxix]['box_coords'] = coords

            # re-transform segmentation predictions.

            chunk_dict['seg_preds'] = np.flip(chunk_dict['seg_preds'], axis=2)

            results_list.append(chunk_dict)

            # second mirroring: x-axis.

            batch['data'] = np.flip(img, axis=3).copy()

            chunk_dict = self.spatial_tiling_forward(batch, mirrored_patch_crops[1], n_aug='2')

            # re-transform coordinates.

            for ix in range(len(chunk_dict['boxes'])):

                for boxix in range(len(chunk_dict['boxes'][ix])):

                    coords = chunk_dict['boxes'][ix][boxix]['box_coords'].copy()

                    coords[1] = org_img_shape[3] - chunk_dict['boxes'][ix][boxix]['box_coords'][3]

                    coords[3] = org_img_shape[3] - chunk_dict['boxes'][ix][boxix]['box_coords'][1]

                    assert coords[2] >= coords[0], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]

                    assert coords[3] >= coords[1], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]

                    chunk_dict['boxes'][ix][boxix]['box_coords'] = coords

            # re-transform segmentation predictions.

            chunk_dict['seg_preds'] = np.flip(chunk_dict['seg_preds'], axis=3)

            results_list.append(chunk_dict)

            # third mirroring: y-axis and x-axis.

            batch['data'] = np.flip(np.flip(img, axis=2), axis=3).copy()

            chunk_dict = self.spatial_tiling_forward(batch, mirrored_patch_crops[2], n_aug='3')

            # re-transform coordinates.

            for ix in range(len(chunk_dict['boxes'])):

                for boxix in range(len(chunk_dict['boxes'][ix])):

                    coords = chunk_dict['boxes'][ix][boxix]['box_coords'].copy()

                    coords[0] = org_img_shape[2] - chunk_dict['boxes'][ix][boxix]['box_coords'][2]

                    coords[2] = org_img_shape[2] - chunk_dict['boxes'][ix][boxix]['box_coords'][0]

                    coords[1] = org_img_shape[3] - chunk_dict['boxes'][ix][boxix]['box_coords'][3]

                    coords[3] = org_img_shape[3] - chunk_dict['boxes'][ix][boxix]['box_coords'][1]

                    assert coords[2] >= coords[0], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]

                    assert coords[3] >= coords[1], [coords, chunk_dict['boxes'][ix][boxix]['box_coords'].copy()]

                    chunk_dict['boxes'][ix][boxix]['box_coords'] = coords

            # re-transform segmentation predictions.

            chunk_dict['seg_preds'] = np.flip(np.flip(chunk_dict['seg_preds'], axis=2), axis=3).copy()

            results_list.append(chunk_dict)

            batch['data'] = img

        # aggregate all boxes/seg_preds per batch element from data_aug predictions.

        results_dict = {}

        results_dict['boxes'] = [[item for d in results_list for item in d['boxes'][batch_instance]]

                                 for batch_instance in range(org_img_shape[0])]

        results_dict['seg_preds'] = np.array([[item for d in results_list for item in d['seg_preds'][batch_instance]]

                                              for batch_instance in range(org_img_shape[0])])

        if self.mode == 'val':

            try:

                results_dict['torch_loss'] = results_list[0]['torch_loss']

                results_dict['class_loss'] = results_list[0]['class_loss']

            except KeyError:

                pass

        return results_dict

    def spatial_tiling_forward(self, batch, patch_crops=None, n_aug='0'):

        """

        forwards batch to batch_tiling_forward method and receives and returns a dictionary with results.

        if patch-based prediction, the results received from batch_tiling_forward will be on a per-patch-basis.

        this method uses the provided patch_crops to re-transform all predictions to whole-image coordinates.

        Patch-origin information of all box-predictions will be needed for consolidation, hence it is stored as

        'patch_id', which is a unique string for each patch (also takes current data aug and temporal epoch instances

        into account). all box predictions get additional information about the amount overlapping patches at the

        respective position (used for consolidation).

        :return. results_dict: stores the results for one patient. dictionary with keys:

                 - 'boxes': list over batch elements. each element is a list over boxes, where each box is

                            one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions,

                            and a dummy batch dimension of 1 for 3D predictions.

                 - 'seg_preds': pixel-wise predictions. (b, 1, y, x, (z))

                 - losses (only in validation mode)

        """

        if patch_crops is not None:

            patches_dict = self.batch_tiling_forward(batch)

            results_dict = {'boxes': [[] for _ in range(batch['original_img_shape'][0])]}

            # instanciate segemntation output array. Will contain averages over patch predictions.

            out_seg_preds = np.zeros(batch['original_img_shape'], dtype=np.float16)[:, 0][:, None]

            # counts patch instances per pixel-position.

            patch_overlap_map = np.zeros_like(out_seg_preds, dtype='uint8')

            #unmold segmentation outputs. loop over patches.

            for pix, pc in enumerate(patch_crops):

                if self.cf.dim == 3:

                    out_seg_preds[:, :, pc[0]:pc[1], pc[2]:pc[3], pc[4]:pc[5]] += patches_dict['seg_preds'][pix][None]

                    patch_overlap_map[:, :, pc[0]:pc[1], pc[2]:pc[3], pc[4]:pc[5]] += 1

                else:

                    out_seg_preds[pc[4]:pc[5], :, pc[0]:pc[1], pc[2]:pc[3], ] += patches_dict['seg_preds'][pix]

                    patch_overlap_map[pc[4]:pc[5], :, pc[0]:pc[1], pc[2]:pc[3], ] += 1

            # take mean in overlapping areas.

            out_seg_preds[patch_overlap_map > 0] /= patch_overlap_map[patch_overlap_map > 0]

            results_dict['seg_preds'] = out_seg_preds

            # unmold box outputs. loop over patches.

            for pix, pc in enumerate(patch_crops):

                patch_boxes = patches_dict['boxes'][pix]

                for box in patch_boxes:

                    # add unique patch id for consolidation of predictions.

                    box['patch_id'] = self.rank_ix + '_' + n_aug + '_' + str(pix)

                    # boxes from the edges of a patch have a lower prediction quality, than the ones at patch-centers.

                    # hence they will be downweighted for consolidation, using the 'box_patch_center_factor', which is

                    # obtained by a normal distribution over positions in the patch and average over spatial dimensions.

                    # Also the info 'box_n_overlaps' is stored for consolidation, which depicts the amount over

                    # overlapping patches at the box's position.

                    c = box['box_coords']

                    box_centers = [(c[ii] + c[ii + 2]) / 2 for ii in range(2)]

                    if self.cf.dim == 3:

                        box_centers.append((c[4] + c[5]) / 2)

                    box['box_patch_center_factor'] = np.mean(

                        [norm.pdf(bc, loc=pc, scale=pc * 0.8) * np.sqrt(2 * np.pi) * pc * 0.8 for bc, pc in

                         zip(box_centers, np.array(self.cf.patch_size) / 2)])

                    if self.cf.dim == 3:

                        c += np.array([pc[0], pc[2], pc[0], pc[2], pc[4], pc[4]])

                        int_c = [int(np.floor(ii)) if ix%2 == 0 else int(np.ceil(ii)) for ix, ii in enumerate(c)]

                        box['box_n_overlaps'] = np.mean(patch_overlap_map[:, :, int_c[1]:int_c[3], int_c[0]:int_c[2], int_c[4]:int_c[5]])

                        results_dict['boxes'][0].append(box)

                    else:

                        c += np.array([pc[0], pc[2], pc[0], pc[2]])

                        int_c = [int(np.floor(ii)) if ix % 2 == 0 else int(np.ceil(ii)) for ix, ii in enumerate(c)]

                        box['box_n_overlaps'] = np.mean(patch_overlap_map[pc[4], :, int_c[1]:int_c[3], int_c[0]:int_c[2]])

                        results_dict['boxes'][pc[4]].append(box)

            if self.mode == 'val':

                try:

                    results_dict['torch_loss'] = patches_dict['torch_loss']

                    results_dict['class_loss'] = patches_dict['class_loss']

                except KeyError:

                    pass

        # if predictions are not patch-based:

        # add patch-origin info to boxes (entire image is the same patch with overlap=1) and return results.

        else:

            results_dict = self.batch_tiling_forward(batch)

            for b in results_dict['boxes']:

                for box in b:

                    box['box_patch_center_factor'] = 1

                    box['box_n_overlaps'] = 1

                    box['patch_id'] = self.rank_ix + '_' + n_aug

        return results_dict

    def batch_tiling_forward(self, batch):

        """

        calls the actual network forward method. in patch-based prediction, the batch dimension might be overladed

        with n_patches >> batch_size, which would exceed gpu memory. In this case, batches are processed in chunks of

        batch_size. validation mode calls the train method to monitor losses (returned ground truth objects are discarded).

        test mode calls the test forward method, no ground truth required / involved.

        :return. results_dict: stores the results for one patient. dictionary with keys:

                 - 'boxes': list over batch elements. each element is a list over boxes, where each box is

                            one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D predictions,

                            and a dummy batch dimension of 1 for 3D predictions.

                 - 'seg_preds': pixel-wise predictions. (b, 1, y, x, (z))

                 - losses (only in validation mode)

        """

        #self.logger.info('forwarding (patched) patient with shape: {}'.format(batch['data'].shape))

        img = batch['data']

        #batch['data'] = torch.from_numpy(batch['data']).float().to(self.device)

        if img.shape[0] <= self.cf.batch_size:

            if self.mode == 'val':

                # call training method to monitor losses

                results_dict = self.net.train_forward(batch, is_validation=True)

                # discard returned ground-truth boxes (also training info boxes).

                results_dict['boxes'] = [[box for box in b if box['box_type'] == 'det'] for b in results_dict['boxes']]

            else:

                results_dict = self.net.test_forward(batch, return_masks=self.cf.return_masks_in_test)

        else:

            split_ixs = np.split(np.arange(img.shape[0]), np.arange(img.shape[0])[::self.cf.batch_size])

            chunk_dicts = []

            for chunk_ixs in split_ixs[1:]:  # first split is elements before 0, so empty

                b = {k: batch[k][chunk_ixs] for k in batch.keys()

                     if (isinstance(batch[k], np.ndarray) and batch[k].shape[0] == img.shape[0])}

                if self.mode == 'val':

                    chunk_dicts += [self.net.train_forward(b, is_validation=True)]

                else:

                    chunk_dicts += [self.net.test_forward(b, return_masks=self.cf.return_masks_in_test)]

            results_dict = {}

            # flatten out batch elements from chunks ([chunk, chunk] -> [b, b, b, b, ...])

            results_dict['boxes'] = [item for d in chunk_dicts for item in d['boxes']]

            results_dict['seg_preds'] = np.array([item for d in chunk_dicts for item in d['seg_preds']])

            if self.mode == 'val':

                try:

                    # estimate metrics by mean over batch_chunks. Most similar to training metrics.

                    results_dict['torch_loss'] = torch.mean(torch.cat([d['torch_loss'] for d in chunk_dicts]))

                    results_dict['class_loss'] = np.mean([d['class_loss'] for d in chunk_dicts])

                except KeyError:

                    # losses are not necessarily monitored

                    pass

                # discard returned ground-truth boxes (also training info boxes).

                results_dict['boxes'] = [[box for box in b if box['box_type'] == 'det'] for b in results_dict['boxes']]

        return results_dict

def apply_wbc_to_patient(inputs):

    """

    wrapper around prediction box consolidation: weighted cluster scoring (wcs). processes a single patient.

    loops over batch elements in patient results (1 in 3D, slices in 2D) and foreground classes,

    aggregates and stores results in new list.

    :return. patient_results_list: list over batch elements. each element is a list over boxes, where each box is

                                 one dictionary: [[box_0, ...], [box_n,...]]. batch elements are slices for 2D

                                 predictions, and a dummy batch dimension of 1 for 3D predictions.

    :return. pid: string. patient id.

    """

    in_patient_results_list, pid, class_dict, wcs_iou, n_ens = inputs

    out_patient_results_list = [[] for _ in range(len(in_patient_results_list))]

    for bix, b in enumerate(in_patient_results_list):

        for cl in list(class_dict.keys()):

            boxes = [(ix, box) for ix, box in enumerate(b) if (box['box_type'] == 'det' and box['box_pred_class_id'] == cl)]

            box_coords = np.array([b[1]['box_coords'] for b in boxes])

            box_scores = np.array([b[1]['box_score'] for b in boxes])

            box_center_factor = np.array([b[1]['box_patch_center_factor'] for b in boxes])

            box_n_overlaps = np.array([b[1]['box_n_overlaps'] for b in boxes])

            box_patch_id = np.array([b[1]['patch_id'] for b in boxes])

            if 0 not in box_scores.shape:

                keep_scores, keep_coords = weighted_box_clustering(

                    np.concatenate((box_coords, box_scores[:, None], box_center_factor[:, None],

                                    box_n_overlaps[:, None]), axis=1), box_patch_id, wcs_iou, n_ens)

                for boxix in range(len(keep_scores)):

                    out_patient_results_list[bix].append({'box_type': 'det', 'box_coords': keep_coords[boxix],

                                             'box_score': keep_scores[boxix], 'box_pred_class_id': cl})

        # add gt boxes back to new output list.

        out_patient_results_list[bix].extend([box for box in b if box['box_type'] == 'gt'])

    return [out_patient_results_list, pid]

def merge_2D_to_3D_preds_per_patient(inputs):

    """

    wrapper around 2Dto3D merging operation. Processes a single patient. Takes 2D patient results (slices in batch dimension)

    and returns 3D patient results (dummy batch dimension of 1). Applies an adaption of Non-Maximum Surpression

    (Detailed methodology is described in nms_2to3D).

    :return. results_dict_boxes: list over batch elements (1 in 3D). each element is a list over boxes, where each box is

                                 one dictionary: [[box_0, ...], [box_n,...]].

    :return. pid: string. patient id.

    """

    in_patient_results_list, pid, class_dict, merge_3D_iou = inputs

    out_patient_results_list = []

    for cl in list(class_dict.keys()):

        boxes, slice_ids = [], []

        # collect box predictions over batch dimension (slices) and store slice info as slice_ids.

        for bix, b in enumerate(in_patient_results_list):

            det_boxes = [(ix, box) for ix, box in enumerate(b) if

                     (box['box_type'] == 'det' and box['box_pred_class_id'] == cl)]

            boxes += det_boxes

            slice_ids += [bix] * len(det_boxes)

        box_coords = np.array([b[1]['box_coords'] for b in boxes])

        box_scores = np.array([b[1]['box_score'] for b in boxes])

        slice_ids = np.array(slice_ids)

        if 0 not in box_scores.shape:

            keep_ix, keep_z = nms_2to3D(

                np.concatenate((box_coords, box_scores[:, None], slice_ids[:, None]), axis=1), merge_3D_iou)

        else:

            keep_ix, keep_z = [], []

        # store kept predictions in new results list and add corresponding z-dimension info to coordinates.

        for kix, kz in zip(keep_ix, keep_z):

            out_patient_results_list.append({'box_type': 'det', 'box_coords': list(box_coords[kix]) + kz,

                                             'box_score': box_scores[kix], 'box_pred_class_id': cl})

    gt_boxes = [box for b in in_patient_results_list for box in b if box['box_type'] == 'gt']

    if len(gt_boxes) > 0:

        assert np.all([len(box["box_coords"]) == 6 for box in gt_boxes]), "expanded preds to 3D but GT is 2D."

    out_patient_results_list += gt_boxes

    # add dummy batch dimension 1 for 3D.

    return [[out_patient_results_list], pid]

def weighted_box_clustering(dets, box_patch_id, thresh, n_ens):

    """

    consolidates overlapping predictions resulting from patch overlaps, test data augmentations and temporal ensembling.

    clusters predictions together with iou > thresh (like in NMS). Output score and coordinate for one cluster are the

    average weighted by individual patch center factors (how trustworthy is this candidate measured by how centered

    its position the patch is) and the size of the corresponding box.

    The number of expected predictions at a position is n_data_aug * n_temp_ens * n_overlaps_at_position

    (1 prediction per unique patch). Missing predictions at a cluster position are defined as the number of unique

    patches in the cluster, which did not contribute any predict any boxes.

    :param dets: (n_dets, (y1, x1, y2, x2, (z1), (z2), scores, box_pc_facts, box_n_ovs)

    :param thresh: threshold for iou_matching.

    :param n_ens: number of models, that are ensembled. (-> number of expected predicitions per position)

    :return: keep_scores: (n_keep)  new scores of boxes to be kept.

    :return: keep_coords: (n_keep, (y1, x1, y2, x2, (z1), (z2)) new coordinates of boxes to be kept.

    """

    dim = 2 if dets.shape[1] == 7 else 3

    y1 = dets[:, 0]

    x1 = dets[:, 1]

    y2 = dets[:, 2]

    x2 = dets[:, 3]

    scores = dets[:, -3]

    box_pc_facts = dets[:, -2]

    box_n_ovs = dets[:, -1]

    areas = (y2 - y1 + 1) * (x2 - x1 + 1)

    if dim == 3:

        z1 = dets[:, 4]

        z2 = dets[:, 5]

        areas *= (z2 - z1 + 1)

    # order is the sorted index.  maps order to index o[1] = 24 (rank1, ix 24)

    order = scores.argsort()[::-1]

    keep = []

    keep_scores = []

    keep_coords = []

    while order.size > 0:

        i = order[0]  # higehst scoring element

        xx1 = np.maximum(x1[i], x1[order])

        yy1 = np.maximum(y1[i], y1[order])

        xx2 = np.minimum(x2[i], x2[order])

        yy2 = np.minimum(y2[i], y2[order])

        w = np.maximum(0.0, xx2 - xx1 + 1)

        h = np.maximum(0.0, yy2 - yy1 + 1)

        inter = w * h

        if dim == 3:

            zz1 = np.maximum(z1[i], z1[order])

            zz2 = np.minimum(z2[i], z2[order])

            d = np.maximum(0.0, zz2 - zz1 + 1)

            inter *= d

        # overall between currently highest scoring box and all boxes.

        ovr = inter / (areas[i] + areas[order] - inter)

        # get all the predictions that match the current box to build one cluster.

        matches = np.argwhere(ovr > thresh)

        match_n_ovs = box_n_ovs[order[matches]]

        match_pc_facts = box_pc_facts[order[matches]]

        match_patch_id = box_patch_id[order[matches]]

        match_ov_facts = ovr[matches]

        match_areas = areas[order[matches]]

        match_scores = scores[order[matches]]

        # weight all socres in cluster by patch factors, and size.

        match_score_weights = match_ov_facts * match_areas * match_pc_facts

        match_scores *= match_score_weights

        # for the weigted average, scores have to be divided by the number of total expected preds at the position

        # of the current cluster. 1 Prediction per patch is expected. therefore, the number of ensembled models is

        # multiplied by the mean overlaps of  patches at this position (boxes of the cluster might partly be

        # in areas of different overlaps).

        n_expected_preds = n_ens * np.mean(match_n_ovs)

        # the number of missing predictions is obtained as the number of patches,

        # which did not contribute any prediction to the current cluster.

        n_missing_preds = np.max((0, n_expected_preds - np.unique(match_patch_id).shape[0]))

        # missing preds are given the mean weighting

        # (expected prediction is the mean over all predictions in cluster).

        denom = np.sum(match_score_weights) + n_missing_preds * np.mean(match_score_weights)

        # compute weighted average score for the cluster

        avg_score = np.sum(match_scores) / denom

        # compute weighted average of coordinates for the cluster. now only take existing

        # predictions into account.

        avg_coords = [np.sum(y1[order[matches]] * match_scores) / np.sum(match_scores),

                      np.sum(x1[order[matches]] * match_scores) / np.sum(match_scores),

                      np.sum(y2[order[matches]] * match_scores) / np.sum(match_scores),

                      np.sum(x2[order[matches]] * match_scores) / np.sum(match_scores)]

        if dim == 3:

            avg_coords.append(np.sum(z1[order[matches]] * match_scores) / np.sum(match_scores))

            avg_coords.append(np.sum(z2[order[matches]] * match_scores) / np.sum(match_scores))

        # some clusters might have very low scores due to high amounts of missing predictions.

        # filter out the with a conservative threshold, to speed up evaluation.

        if avg_score > 0.01:

            keep_scores.append(avg_score)

            keep_coords.append(avg_coords)

        # get index of all elements that were not matched and discard all others.

        inds = np.where(ovr <= thresh)[0]

        order = order[inds]

    return keep_scores, keep_coords

def nms_2to3D(dets, thresh):

    """

    Merges 2D boxes to 3D cubes. Therefore, boxes of all slices are projected into one slices. An adaptation of Non-maximum surpression

    is applied, where clusters are found (like in NMS) with an extra constrained, that surpressed boxes have to have 'connected'

    z-coordinates w.r.t the core slice (cluster center, highest scoring box). 'connected' z-coordinates are determined

    as the z-coordinates with predictions until the first coordinate, where no prediction was found.

    example: a cluster of predictions was found overlap > iou thresh in xy (like NMS). The z-coordinate of the highest

    scoring box is 50. Other predictions have 23, 46, 48, 49, 51, 52, 53, 56, 57.

    Only the coordinates connected with 50 are clustered to one cube: 48, 49, 51, 52, 53. (46 not because nothing was

    found in 47, so 47 is a 'hole', which interrupts the connection). Only the boxes corresponding to these coordinates

    are surpressed. All others are kept for building of further clusters.

    This algorithm works better with a certain min_confidence of predictions, because low confidence (e.g. noisy/cluttery)

    predictions can break the relatively strong assumption of defining cubes' z-boundaries at the first 'hole' in the cluster.

    :param dets: (n_detections, (y1, x1, y2, x2, scores, slice_id)

    :param thresh: iou matchin threshold (like in NMS).

    :return: keep: (n_keep) 1D tensor of indices to be kept.

    :return: keep_z: (n_keep, [z1, z2]) z-coordinates to be added to boxes, which are kept in order to form cubes.

    """

    y1 = dets[:, 0]

    x1 = dets[:, 1]

    y2 = dets[:, 2]

    x2 = dets[:, 3]

    scores = dets[:, -2]

    slice_id = dets[:, -1]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)

    order = scores.argsort()[::-1]

    keep = []

    keep_z = []

    while order.size > 0:  # order is the sorted index.  maps order to index o[1] = 24 (rank1, ix 24)

        i = order[0]  # pop higehst scoring element

        xx1 = np.maximum(x1[i], x1[order])

        yy1 = np.maximum(y1[i], y1[order])

        xx2 = np.minimum(x2[i], x2[order])

        yy2 = np.minimum(y2[i], y2[order])

        w = np.maximum(0.0, xx2 - xx1 + 1)

        h = np.maximum(0.0, yy2 - yy1 + 1)

        inter = w * h

        ovr = inter / (areas[i] + areas[order] - inter)

        matches = np.argwhere(ovr > thresh)  # get all the elements that match the current box and have a lower score

        slice_ids = slice_id[order[matches]]

        core_slice = slice_id[int(i)]

        upper_wholes = [ii for ii in np.arange(core_slice, np.max(slice_ids)) if ii not in slice_ids]

        lower_wholes = [ii for ii in np.arange(np.min(slice_ids), core_slice) if ii not in slice_ids]

        max_valid_slice_id = np.min(upper_wholes) if len(upper_wholes) > 0 else np.max(slice_ids)

        min_valid_slice_id = np.max(lower_wholes) if len(lower_wholes) > 0 else np.min(slice_ids)

        z_matches = matches[(slice_ids <= max_valid_slice_id) & (slice_ids >= min_valid_slice_id)]

        z1 = np.min(slice_id[order[z_matches]]) - 1

        z2 = np.max(slice_id[order[z_matches]]) + 1

        keep.append(i)

        keep_z.append([z1, z2])

        order = np.delete(order, z_matches, axis=0)

    return keep, keep_z

def get_mirrored_patch_crops(patch_crops, org_img_shape):

    """

    apply 3 mirrror transformations (x-axis, y-axis, x&y-axis)

    to given patch crop coordinates and return the transformed coordinates.

    Handles 2D and 3D coordinates.

    :param patch_crops: list of crops: each element is a list of coordinates for given crop [[y1, x1, ...], [y1, x1, ..]]

    :param org_img_shape: shape of patient volume used as world coordinates.

    :return: list of mirrored patch crops: lenght=3. each element is a list of transformed patch crops.

    """

    mirrored_patch_crops = []

    # y-axis transform.

    mirrored_patch_crops.append([[org_img_shape[2] - ii[1],

                                  org_img_shape[2] - ii[0],

                                  ii[2], ii[3]] if len(ii) == 4 else

                                 [org_img_shape[2] - ii[1],

                                  org_img_shape[2] - ii[0],

                                  ii[2], ii[3], ii[4], ii[5]] for ii in patch_crops])

    # x-axis transform.

    mirrored_patch_crops.append([[ii[0], ii[1],

                                  org_img_shape[3] - ii[3],

                                  org_img_shape[3] - ii[2]] if len(ii) == 4 else

                                 [ii[0], ii[1],

                                  org_img_shape[3] - ii[3],

                                  org_img_shape[3] - ii[2],

                                  ii[4], ii[5]] for ii in patch_crops])

    # y-axis and x-axis transform.

    mirrored_patch_crops.append([[org_img_shape[2] - ii[1],

                                  org_img_shape[2] - ii[0],

                                  org_img_shape[3] - ii[3],

                                  org_img_shape[3] - ii[2]] if len(ii) == 4 else

                                 [org_img_shape[2] - ii[1],

                                  org_img_shape[2] - ii[0],

                                  org_img_shape[3] - ii[3],

                                  org_img_shape[3] - ii[2],

                                  ii[4], ii[5]] for ii in patch_crops])

    return mirrored_patch_crops