"""
\file ThresholdMaskingScheme.py
\brief      Class to generate and evaluate masks obtained by simple thresholding
\author     Michael Ebner (michael.ebner.14@ucl.ac.uk)
\date       June 2017
"""

import numpy as np
import src.utilities as utils
import src.Exceptions as Exceptions
from src.MaskingScheme import MaskingScheme


class ThresholdMaskingScheme(MaskingScheme):
    """
    Estimate the blood pool mask by simple thresholding given a target and 
    image array according to the format defined by the class DataBase.

    \details    It is assumed that the target array is being created based on 
                i- and o-contour information. The region surrounded by the 
                i-contours defines the blood pool (label=2) and the region 
                between o- and i-contours corresponds to the heart muscle 
                (label=1), i.e.

                data_blood_pool 
                    = images_array[np.where(targets_array == 2)]

                data_heart_muscle
                    = images_array[np.where(targets_array == 1)]

    """

    def __init__(self, images_array=None, targets_array=None,
                 label_ground_truth_target=2, thresholds_list=[0, 250]):
        """
        Store information on given image and target data array where the target
        array includes labelling information.

        \param      images_array               Images data as numpy array
        \param      targets_array              Targets data as numpy array
        \param      label_ground_truth_target  index referring to the label in
                                               targets_array being considered
                                               as "ground-truth"
        \param      thresholds_list            list of thresholds for automatic
                                               sweeping
        """
        MaskingScheme.__init__(self, images_array=images_array, targets_array=targets_array)
        
        self._label_ground_truth_target = label_ground_truth_target
        self._thresholds_list = thresholds_list

    def set_thresholds(self, thresholds_list):
        """!
        Set list of thresholds
        
        \param      thresholds_list  List of threshold parameters
        """
        self._thresholds_list = thresholds_list

    def estimate_optimal_parameter(self):
        """!
        Estimate optimal threshold with respect to Dice score by sweeping 
        through all provided thresholds.

        \return     Optimal threshold corresponding to highest dice score
        """

        dice_scores_mean = self.evaluate_masking_scheme_by_threshold_sweeping()

        optimal_threshold = self._thresholds_list[np.argmax(dice_scores_mean)]

        return optimal_threshold

    def evaluate_masking_scheme_by_threshold_sweeping(self):
        """!
        Sweep through all specified threshold values and return the mean dice
        scores.

        \details    For each threshold the mean dice score for the stored pair
                    of images and target slices are computed. The rationale
                    is to find the "optimal" threshold value afterwards.

        \return     Mean dice scores as list of length equal to the number of
                    specified thresholds to sweep through
        """
        dice_scores_mean = [self.get_mean_dice_score(self._thresholds_list[i]).mean()
                            for i in range(0, len(self._thresholds_list))]

        return dice_scores_mean

    def get_mean_dice_score(self, threshold):
        """!
        Compute mean dice score given a specified threshold

        \param      threshold  threshold as integer value

        \return     The mean dice score.
        """

        # Estimate targets/masks given the set threshold
        targets_array_estimate = self.get_target_array_estimate(threshold)

        # Get ground-truth targets/masks of blood pool
        targets_array_ground_truth = self._get_ground_truth_targets_array()

        # Get Dice scores for each single slice
        dice_scores = [utils.dice_score(targets_array_ground_truth[:, :, j],
                                        targets_array_estimate[:, :, j]) for j in range(0, targets_array_estimate.shape[2])]

        # Return mean dices score of all slices
        return np.mean(dice_scores)

    def get_target_array_estimate(self, threshold):
        """!
        Gets the target array estimate, i.e. mask of blood pool, by applying 
        a simple thresholding.

        \details    A target mask is being generated by finding the values
                    within the o-contours region greater than the threshold

        \param      threshold  The threshold

        \return     The target array estimate as numpy array
        """

        self._check_input()

        # Recover o-contours given the labels in targets array
        mask_ocontours = np.zeros(self._images_array.shape, dtype=bool)
        mask_ocontours[np.where(self._targets_array > 0)] = True

        # Within the contours apply a thresholding to estimate masks
        target_array_estimate = np.zeros_like(mask_ocontours)
        images_array_masked = self._images_array * mask_ocontours
        target_array_estimate[
            np.where(images_array_masked > threshold)] = True

        # utils.show_image_data(
        #     self._images_array,
        #     target_array_estimate,
        #     title="icontours_est")

        return target_array_estimate

    def _get_ground_truth_targets_array(self):

        targets_array_ground_truth = np.zeros_like(self._targets_array)
        targets_array_ground_truth[
            np.where(self._targets_array == self._label_ground_truth_target)] = True

        return targets_array_ground_truth

    def _check_input(self):
        """!
        Check given input variables
        """

        if self._images_array is None:
            raise Exceptions.ObjectNotCreated("set_images_array")

        if self._targets_array is None:
            raise Exceptions.ObjectNotCreated("set_targets_array")

        if type(self._thresholds_list) is not list:
            raise Exceptions.ObjectIsNotList()
