# Copyright 2019 Diamond Light Source Ltd.
#
# 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.

"""
.. module:: A plugin to initialise a binary mask for level sets and distance transform segmentations
   :platform: Unix
   :synopsis: A plugin to initialise a binary mask for level sets and distance transform segmentations

.. moduleauthor:: Daniil Kazantsev <scientificsoftware@diamond.ac.uk>
"""

from savu.plugins.plugin import Plugin
from savu.plugins.driver.cpu_plugin import CpuPlugin
from savu.plugins.utils import register_plugin

import numpy as np

# using Morphological snakes module from
# https://github.com/pmneila/morphsnakes
from morphsnakes import circle_level_set

@register_plugin
class MaskInitialiser(Plugin, CpuPlugin):
    """
    A plugin to initialise a binary mask for level sets and distance transform segmentations.\
    Seeds are generated by providing coordinates of three points in 3D space (start-middle-finish) \
    and connecting them with a cylinder of a certain radius. \
    Importantly the Z coordinate is given following VOLUME_XY vertical pattern

    :param mask1_coordinates: X0,Y0,Z0 (start) X1,Y1,Z1 (middle) and X2,Y2,Z2 (finish) coordinates of three points. Default: [10, 10, 0, 15, 15, 15, 20, 20, 20].
    :param mask1_radius: Mask1 will be initialised with an ellipse of radius. Default: 5.
    :param mask2_coordinates: The second mask coordinates. Default: None.
    :param mask2_radius: Mask2 will be initialised with an ellipse of radius. Default: None.
    :param out_datasets: The default names . Default: ['INIT_MASK'].
    """

    def __init__(self):
        super(MaskInitialiser, self).__init__("MaskInitialiser")

    def setup(self):
        in_dataset, out_dataset = self.get_datasets()
        in_pData, out_pData = self.get_plugin_datasets()
        in_pData[0].plugin_data_setup('VOLUME_YZ', 'single')

        out_dataset[0].create_dataset(in_dataset[0], dtype=np.uint8)
        out_pData[0].plugin_data_setup('VOLUME_YZ', 'single')
        self.full_input_shape = in_dataset[0].get_shape()

    def pre_process(self):
        # extract given parameters
        self.mask1_coordinates = self.parameters['mask1_coordinates']
        self.mask1_radius = self.parameters['mask1_radius']
        self.mask2_coordinates = self.parameters['mask2_coordinates']
        self.mask2_radius = self.parameters['mask2_radius']

        self.coordX = [0.0, 0.0, 0.0]
        self.coordY = [0.0, 0.0, 0.0]
        self.coordZ = [0.0, 0.0, 0.0]

        self.coordX[0] = self.mask1_coordinates[0]
        self.coordY[0] = self.mask1_coordinates[1]
        self.coordZ[0] = self.mask1_coordinates[2]
        self.coordX[1] = self.mask1_coordinates[3]
        self.coordY[1] = self.mask1_coordinates[4]
        self.coordZ[1] = self.mask1_coordinates[5]
        self.coordX[2] = self.mask1_coordinates[6]
        self.coordY[2] = self.mask1_coordinates[7]
        self.coordZ[2] = self.mask1_coordinates[8]

        if (self.mask2_coordinates is not None):
            self.coord2X = [0.0, 0.0, 0.0]
            self.coord2Y = [0.0, 0.0, 0.0]
            self.coord2Z = [0.0, 0.0, 0.0]

            self.coord2X[0] = self.mask2_coordinates[0]
            self.coord2Y[0] = self.mask2_coordinates[1]
            self.coord2Z[0] = self.mask2_coordinates[2]
            self.coord2X[1] = self.mask2_coordinates[3]
            self.coord2Y[1] = self.mask2_coordinates[4]
            self.coord2Z[1] = self.mask2_coordinates[5]
            self.coord2X[2] = self.mask2_coordinates[6]
            self.coord2Y[2] = self.mask2_coordinates[7]
            self.coord2Z[2] = self.mask2_coordinates[8]

    def process_frames(self, data):
        # get the index of a current frame
        index_current = self.get_plugin_in_datasets()[0].get_current_frame_idx()
        [dimX,dimY] = np.shape(data[0])
        mask = np.uint8(np.zeros(np.shape(data[0])))
        mask = mask_gen(mask, self.coordX, self.coordY, self.coordZ, self.mask1_radius, dimX, dimY, index_current)
        if ((self.mask2_coordinates is not None) and (self.mask2_radius is not None)):
            mask2 = np.uint8(np.zeros(np.shape(data[0])))
            mask2 = mask_gen(mask2, self.coord2X, self.coord2Y, self.coord2Z, self.mask2_radius, dimX, dimY, index_current)
            mask = np.add(mask, mask2)
        return mask

    def nInput_datasets(self):
        return 1
    def nOutput_datasets(self):
        return 1
    def get_max_frames(self):
        return 'single'


def mask_gen(mask, coordX, coordY, coordZ, mask_radius, dimX, dimY, index_current):
    steps1 = coordZ[1] - coordZ[0]
    distance1 = np.sqrt((coordX[1] - coordX[0])**2 + (coordY[1] - coordY[0])**2)
    d_dist1 = distance1/(steps1 - 1.0)
    d_step1 = d_dist1
    steps2 = coordZ[2] - coordZ[1]
    distance2 = np.sqrt((coordX[2] - coordX[1])**2 + (coordY[2] - coordY[1])**2)
    d_dist2 = distance2/(steps2 - 1.0)
    d_step2 = d_dist2
    if ((index_current >= coordZ[0]) & (index_current <= coordZ[1])):
        if ((coordX[0] == 0) & (coordY[0] == 0) & (coordX[1] == 0) & (coordY[1] == 0)):
            # create a full region mask (except the boundaries)
            bound_width=2 # outer boundary width
            mask[bound_width:dimX-bound_width, bound_width:dimY-bound_width] = 1
        else:
            d_step1 = (index_current - coordZ[0])*d_dist1
            if (distance1 != 0.0):
                t = d_step1/distance1
            else:
                t = 0.0
            if (coordX[0] == coordX[1]):
                x_t = np.int(coordX[0])
            else:
                x_t1 = np.round((1.0 - t)*coordX[0] + t*coordX[1])
                x_t = np.int(x_t1[0])
            if(coordY[0] == coordY[1]):
                y_t = np.int(coordY[0])
            else:
                y_t1 = np.round((1.0 - t)*coordY[0] + t*coordY[1])
                y_t = np.int(y_t1[0])
            mask = np.uint8(circle_level_set(np.shape(mask), (y_t, x_t), mask_radius))
    if ((index_current > coordZ[1]) & (index_current <= coordZ[2])):
        if ((coordX[1] == 0) & (coordY[1] == 0) & (coordX[2] == 0) & (coordY[2] == 0)):
            # create a full region mask (except the boundaries)
            bound_width=2 # outer boundary width
            mask[bound_width:dimX-bound_width, bound_width:dimY-bound_width] = 1
        else:
            d_step2 = (index_current - coordZ[1])*d_dist2
            if (distance2 != 0.0):
                t = d_step2/distance2
            else:
                t = 0.0
            if (coordX[1] == coordX[2]):
                x_t = np.int(coordX[1])
            else:
                x_t1 = np.round((1.0 - t)*coordX[1] + t*coordX[2])
                x_t = np.int(x_t1[0])
            if(coordY[1] == coordY[2]):
                y_t = np.int(coordY[1])
            else:
                y_t1 = np.round((1.0 - t)*coordY[1] + t*coordY[2])
                y_t = np.int(y_t1[0])
            mask = np.uint8(circle_level_set(np.shape(mask), (y_t, x_t), mask_radius))
    return [mask]
