"""

Inference module supporting whole slide images using TF Serving.

Eric Wang

Email: eric.wang@thorough.ai

"""

import os

import shutil

import numpy as np

from utils.tf_serving import TFServing

from utils.slide import Slide

from utils import config

from utils.libs import write, generate_effective_regions, generate_overlap_tile, \

    post_processing, concat_patches

class Inference:

    def __init__(self, data_dir, data_list, class_num, result_dir, use_level):

        """

        This is the main inference module for the sake of easy to call.

        :param data_dir: The directory storing the while image slides.

        :param data_list: The text file indicating the slide names.

        :param class_num: Number of predicted classes.

        :param result_dir: Where to put the predicted results.

        :param use_level: Which slide size we want to analyze, 0 for 40x, 1 for 20x, etc.

        """

        if data_dir.endswith('/'):

            self.data_dir = data_dir

        else:

            self.data_dir = data_dir + '/'

        self.data_list = data_list

        self.class_num = class_num

        if result_dir.endswith('/'):

            self.result_dir = result_dir

        else:

            self.result_dir = result_dir + '/'

        if not os.path.exists(self.result_dir):

            os.mkdir(self.result_dir)

        self.use_level = use_level

        self.config = config

    @staticmethod

    def _infer(tfs_client, image):

        """

        Inference for an image patch using TF Serving.

        :param tfs_client: TF Serving client.

        :param image: The image patch.

        :return: Predicted heatmap.

        """

        try:

            prediction = tfs_client.predict(image, config.MODEL_NAME)

        except Exception as e:

            print('TF_SERVING_HOST: {}'.format(config.TF_SERVING_HOST))

            print(e)

            raise

        else:

            return prediction

    def run(self):

        """

        Proceeds the inference procedure.

        """

        inference_list = open(self.data_list).readlines()

        tfs_client = TFServing(config.TF_SERVING_HOST, config.TF_SERVING_PORT)

        for item in inference_list:

            image_name, image_suffix = item.split('\n')[0].split('/')[-1].split('.')

            print('[INFO] Analyzing: ' + self.data_dir + item.split('\n')[0])

            if not image_suffix in self.config.format_mapping.keys():

                print('[ERROR] File ' + item + ' format not supported yet.')

                continue

            image_handle = Slide(self.data_dir + item.split('\n')[0])

            image_dimensions = image_handle.level_dimensions[self.use_level]

            regions = generate_effective_regions(image_dimensions)

            index = 0

            region_num = len(regions)

            temp_dir = self.config.TEMP_DIR + image_name + '/'

            if not os.path.exists(temp_dir):

                os.makedirs(temp_dir)

            for region in regions:

                shifted_region, clip_region = generate_overlap_tile(region, image_dimensions)

                index += 1

                if index % 1 == 0:

                    print('[INFO]  Progress: ' + str(index) + ' / ' + str(region_num))

                input_image = np.array(image_handle.read_region(

                    location=(int(shifted_region[0]),

                              int(shifted_region[1])),

                    level=self.use_level, size=(self.config.PATCH_SIZE, self.config.PATCH_SIZE)))[:, :, 0: 3]

                prediction_result = self._infer(tfs_client, input_image)

                prediction_result = prediction_result[clip_region[0]: (self.config.CENTER_SIZE + clip_region[0]),

                                    clip_region[1]: (self.config.CENTER_SIZE + clip_region[1])]

                prediction_result = prediction_result[region[2]:(region[4] + 1), region[3]:(region[5] + 1)]

                if self.config.DO_POST_PROCESSING:

                    prediction_result = post_processing(prediction_result)

                write(temp_dir + image_name + '_' + str(region[0]) + '_' + str(region[1])

                      + '_prediction.png', prediction_result, self.class_num)

            print('[INFO] Postprocessing...')

            full_prediction = concat_patches(temp_dir, image_name)

            write(self.result_dir +

                  '_'.join([image_name, 'prediction_thumbnail']) + '.png', full_prediction, color_map=False)

            if not self.config.KEEP_TEMP:

                shutil.rmtree(temp_dir)

            print('[INFO] Prediction saved to ' + self.result_dir + '_'.join(

                [image_name, 'prediction_thumbnail']) + '.png')