'''Classes and methods for segmentation of reticular objects within cells.'''


## IMPORT DEPENDENCIES ##
import matplotlib
matplotlib.use('Agg')
import os
import sys
import pickle
import time
from operator import itemgetter
import numpy as np
import pandas as pd
from skimage import io
from skimage.morphology import watershed, disk
from skimage.feature import canny
from scipy.ndimage.filters import gaussian_filter, maximum_filter
from scipy.ndimage.morphology import generate_binary_structure, binary_closing
from scipy.ndimage.morphology import distance_transform_edt
from scipy.ndimage.morphology import binary_erosion, binary_dilation
from scipy.ndimage.morphology import binary_fill_holes, binary_opening
from scipy.ndimage import generic_gradient_magnitude, sobel
import matplotlib.pyplot as plt


class MitoSegmentObj:
    '''A container class for objects generated by MitoSegmenter.segment().

    Objects of class MitoSegmentObj contain a raw multipage TIFF image in a
    numpy ndarray format, a similarly formatted image with the segmentation
    output from MitoSegmenter.segment(), a number of segmentation intermediates
    that may be useful for diagnosis of segmentation problems, and relevant
    metadata that may be useful for analysis. Class methods for output of image
    data in TIFF format, for saving of object data in csv format, and for
    pickling the entire object are included.

    IMPORTANT: Do not call this class on its own. Objects of this class are
    generated as output from MitoSegmenter.segment(), which provides all of the
    parameters described in __init__.

    Args !!!IMPORTANT: DO NOT CALL! Passed by MitoSegmenter.segment()!!!
        f_directory (str): Path to the raw image used for segmentation.
        filename (str): raw image filename.
        raw_img (np.ndarray of ints): pixel intensity values of raw input image
            used as a starting point for segmentation. Each array position 
            represents a single pixel.
        gaussian_img (np.ndarray of ints): pixel intensity values of gaussian
            filter output from MitoSegmenter.segment().
        seg_method (str): Segmentation method provided to
            MitoSegmenter.__init__().
        mode (str): Segmentation mode provided to MitoSegmenter.__init__().
        threshold_img (binary np.ndarray): Binary array of pixels corresponding 
            to a segmented object.
        dist_map (np.ndarray of ints): ndarray of the Euclidean distance of
            each pixel marked as 1 in threshold_img to a background pixel (a
            0). Generated by MitoSegmenter.segment().
        smooth_dist_map (np.ndarray of ints): Smoothed distance map generated
            by MitoSegmenter.segment().
        maxima (binary np.ndarray): Local maxima from the smooth_dist_map,
            generated by MitoSegmenter.segment(). 
        labs (np.ndarray of ints): Labeled starting points for watershed
            segmentation. Generated by MitoSegmenter.segment().
        watershed_output (np.ndarray of ints): Segmented objects, the final output
            from MitoSegmenter.segment(). Each array position represents a
            single pixel, and all indices with a given integer value make up
            one segmented object. 0 represents background.
        obj_nums (list of ints): The list of numerical values assigned to 
            segmented objects.
        volumes (dict): A dictionary with obj_num:volume pairs, with volume
            being the size of the segmented object in pixels.
        to_pdout (list of strings): The names of the attributes to be passed to
            to_csv() and to_pandas() for output. Assigned by
            MitoSegmenter.segment().
        mode_params (dict): Dict of additional variables from
            MitoSegmenter.segment() passed to MitoSegmentObj for inclusion in
            attributes. Varies depending upon segmentation method. See
            MitoSegmenter.segment() code for details.

    Attributes:    
        f_directory (str): Path to the raw image used for segmentation.
        filename (str): raw image filename.
        raw_img (np.ndarray of ints): pixel intensity values of raw input image
            used as a starting point for segmentation. Each array position 
            represents a single pixel.
        gaussian_img (np.ndarray of ints): pixel intensity values of gaussian
            filter output from MitoSegmenter.segment().
        seg_method (str): Segmentation method provided to
            MitoSegmenter.__init__().
        mode (str): Segmentation mode provided to MitoSegmenter.__init__().
        threshold_img (binary np.ndarray): Binary array of pixels corresponding 
            to a segmented object.
        dist_map (np.ndarray of ints): ndarray of the Euclidean distance of
            each pixel marked as 1 in threshold_img to a background pixel (a
            0). Generated by MitoSegmenter.segment().
        smooth_dist_map (np.ndarray of ints): Smoothed distance map generated
            by MitoSegmenter.segment().
        maxima (binary np.ndarray): Local maxima from the smooth_dist_map,
            generated by MitoSegmenter.segment(). 
        labs (np.ndarray of ints): Labeled starting points for watershed
            segmentation. Generated by MitoSegmenter.segment().
        mitochondria (np.ndarray of ints): Segmented objects, the final output
            from MitoSegmenter.segment(). Each array position represents a
            single pixel, and all indices with a given integer value make up
            one segmented object. 0 represents background.
        slices (int): The number of Z-slices in the image.
        height (int): The Y-direction image size in pixels.
        width (int): The X-direction image size in pixels.
        obj_nums (list of ints): The list of numerical values assigned to 
            segmented objects.
        nmitos (int): The number of segmented objects.
        volumes (dict): A dictionary with obj_num:volume pairs, with volume
            being the size of the segmented object in pixels.
        volumes_flag (str): The units for volume measurement. Always assigned
            as 'pixels'.
        pdout (list of strings): The names of the attributes to be passed to
            to_csv() and to_pandas() for output. Assigned by
            MitoSegmenter.segment().
        border_rm_flag (bool): Indication of whether or not objects on the edge
            of the image are removed. Defaults to False.
    '''
    def __init__(self, f_directory, filename, raw_img, gaussian_img, 
                 seg_method, mode, threshold_img, dist_map,
                 smooth_dist_map, maxima, labs, watershed_output, 
                 obj_nums, volumes, to_pdout = [], 
                 mode_params = {}):
        '''Initialize the MitoSegmentObj with segmentation data.'''
        print('creating MitoSegmentObj...')
        self.f_directory = f_directory
        self.filename = os.path.basename(filename).lower()
        self.raw_img = raw_img.astype('uint16')
        self.gaussian_img = gaussian_img.astype('uint16')
        self.seg_method = seg_method
        self.mode = mode
        self.threshold_img = threshold_img.astype('uint16')
        self.dist_map = dist_map.astype('uint16')
        self.smooth_dist_map = smooth_dist_map.astype('uint16')
        self.maxima = maxima.astype('uint16')
        self.labs = labs.astype('uint16')
        self.mitochondria = watershed_output.astype('uint16')
        self.slices = self.raw_img.shape[0]
        self.height = self.raw_img.shape[1]
        self.width = self.raw_img.shape[2]
        self.obj_nums = obj_nums
        self.nmitos = len(self.obj_nums)
        self.volumes = volumes
        self.volumes_flag = 'pixels'
        self.pdout = []
        self.border_rm_flag = False
        for key in mode_params:
        # raise an error if an attribute is somehow passed twice
            if hasattr(self, key):
                raise AttributeError('Two copies of the attribute ' + key + 
                                     'were provided to MitoSegmentObj.__init__()')
                setattr(self, key, mode_params[key])
        if to_pdout != []:
            for x in to_pdout:
                self.pdout.append(x)


    def __repr__(self):
        return 'MitoSegmentObj '+ self.filename

    def plot_raw_img(self, display = False):
        '''Plot the raw image using matplotlib.'''
        self.plot_stack(self.raw_img, colormap = 'gray')
        if display == True:
            plt.show()
    def plot_gaussian_img(self, display = False):
        '''Plot the gaussian image using matplotlib.'''
        self.plot_stack(self.gaussian_img, colormap = 'gray')
        if display == True:
            plt.show()
    def plot_threshold_img(self, display = False):
        '''Plot the threshold image using matplotlib.'''
        self.plot_stack(self.threshold_img, colormap = 'gray')
        if display == True:
            plt.show()
    def plot_dist_map(self, display = False):
        '''Plot the distance map image using matplotlib.'''
        self.plot_stack(self.dist_map)
        if display == True:
            plt.show()
    def plot_smooth_dist_map(self, display = False):
        '''Plot the smoothed distance map image using matplotlib.'''
        self.plot_stack(self.smooth_dist_map)
        if display == True:
            plt.show()
    def plot_maxima(self, display = False):
        '''Plot the maxima image using matplotlib.'''
        # expand maxima to make them more easily visible in the output.
        vis_maxima = binary_dilation(self.maxima,
                                     structure = np.ones(shape = (1,5,5)))
        masked_maxima = np.ma.masked_where(vis_maxima == 0, vis_maxima)
        self.plot_maxima_stack(masked_maxima, self.smooth_dist_map)
        if display == True:
            plt.show()
    def plot_watershed(self, display = False):
         '''Plot the segmented objects image using matplotlib.'''
         self.plot_stack(self.mitochondria)
         if display == True:
             plt.show()

    def output_all_images(self, output_dir = None):
        '''Write all images to a new directory.
        
        Write all images associated with the MitoSegmentObj to a new directory. 
        Name that directory according to the filename of the initial image that
        the object was derived from, unless an output directory is provided.
        
        Args:
            output_dir (str, optional): The directory to output image files to.
                If not provided, the images will be a subdirectory to the
                directory containing the initial raw image, which is named based
                on the filename of the raw image.
        '''
        if output_dir == None:
            # name as described in args
            output_dir = self.f_directory + '/' + self.filename[0:self.filename.index('.tif')]
        # make the output directory if it doesn't already exist
        if not os.path.isdir(output_dir):
            print('creating output directory...')
            os.mkdir(output_dir)
        os.chdir(output_dir)
        print('writing images...')
        # save numpy ndarrays as tif images using skimage.io.imsave
        io.imsave('raw_'+self.filename, self.raw_img)
        io.imsave('gaussian_'+self.filename, self.gaussian_img)
        io.imsave('threshold_'+self.filename, self.threshold_img)
        io.imsave('dist_'+self.filename, self.dist_map)
        io.imsave('smooth_dist_'+self.filename,self.smooth_dist_map)
        io.imsave('maxima_'+self.filename,self.maxima)
        io.imsave('wshed_'+self.filename,self.mitochondria)
        if hasattr(self,'edges'):
            io.imsave('edges_'+self.filename,self.edges)

    def output_image(self, imageattr, output_dir = None):
        '''Write one specific image attribute to a new directory.
        
        Write an image associated with the MitoSegmentObj to a new directory. 
        Name that directory according to the filename of the initial image that
        the object was derived from, unless an output directory is provided.
        
        Args:
            imageattr (str): The name of the image attribute to be saved.
            output_dir (str, optional): The directory to output image files to.
                If not provided, the images will be a subdirectory to the
                directory containing the initial raw image, which is named based
                on the filename of the raw image.
        '''
        if output_dir == None:
            output_dir = self.f_directory + '/' + self.filename[0:self.filename.index('.tif')]
        # if the output directory doesn't already exist, make it.
        if not os.path.isdir(output_dir):
            print('creating output directory...')
            os.mkdir(output_dir)
        os.chdir(output_dir)
        print('writing image' + str(imageattr))
        # save the ndarray as a tif image using skimage.io.imsave
        io.imsave(str(imageattr)+self.filename, getattr(self,str(imageattr)))

    def output_plots(self):
        '''Write PDFs of slice-by-slice plots.
        
        Output: PDF plots of each image within MitoSegmentObj in a directory
            named for the original filename they were generated from. Plots are
            generated using the plot_stack method and plotting methods defined
            here.
        '''
        os.chdir(self.f_directory)
        if not os.path.isdir(self.f_directory + '/' +
                             self.filename[0:self.filename.index('.tif')]):
            print('creating output directory...')
            os.mkdir(self.f_directory + '/' +
                     self.filename[0:self.filename.index('.tif')])
        os.chdir(self.f_directory + '/' +
                 self.filename[0:self.filename.index('.tif')])
        print('saving plots...')
        self.plot_raw_img()
        plt.savefig('praw_'+self.filename[0:self.filename.index('.tif')]+'.pdf')
        self.plot_gaussian_img()
        plt.savefig('pgaussian_' +
                    self.filename[0:self.filename.index('.tif')]+'.pdf')
        self.plot_threshold_img()
        plt.savefig('pthreshold_' +
                    self.filename[0:self.filename.index('.tif')]+'.pdf')
        plt.savefig('pdist_' +
                    self.filename[0:self.filename.index('.tif')]+'.pdf')
        self.plot_smooth_dist_map()
        plt.savefig('psmooth_dist_' + 
                    self.filename[0:self.filename.index('.tif')]+'.pdf')
        self.plot_maxima()
        plt.savefig('pmaxima_' +
                    self.filename[0:self.filename.index('.tif')]+'.pdf')
        self.plot_watershed()
        plt.savefig('pwshed_' +
                    self.filename[0:self.filename.index('.tif')]+'.pdf')
    def pickle(self, output_dir = None, filename = None):
        '''pickle the MitoSegmentObj for later loading.'''
        if output_dir == None:
            output_dir = self.f_directory + '/' + self.filename[0:self.filename.index('.tif')]
        if filename == None:
            filename = self.filename[0:self.filename.index('.tif')] + '.pickle'
        if not os.path.isdir(output_dir):
            print('creating output directory...')
            os.mkdir(output_dir)
        os.chdir(output_dir)
        with open(filename, 'wb') as f:
            pickle.dump(self, f, pickle.HIGHEST_PROTOCOL)
        f.close()
    def output_all(self):
        '''Output images as tifs and plots as pdfs, and pickle object.'''
        os.chdir(self.f_directory)
        # make output directories if they don't exist
        if not os.path.isdir(self.f_directory + '/' +
                             self.filename[0:self.filename.index('.tif')]):
            os.mkdir(self.f_directory + '/' +
                     self.filename[0:self.filename.index('.tif')])
        os.chdir(self.f_directory + '/' +
                 self.filename[0:self.filename.index('.tif')])
        print('outputting all data...')
        self.output_plots()
        self.output_all_images()
        self.pickle()
        # TODO: UPDATE THIS METHOD TO INCLUDE PANDAS OUTPUT
    def to_csv(self, output_dir = None):
        '''Output attributes designated by self.pdout in csv format.
        
        Args:
            output_dir (str, optional): Directory to output csv to. If not
            provided, a new subdirectory is created within the directory
            containing the raw image named based on the raw image name, and the
            csv is saved there.
        '''
        os.chdir(self.f_directory)
        if output_dir == None:
            output_dir = self.f_directory + '/' + self.filename[0:self.filename.index('.tif')]
        if not os.path.isdir(output_dir):
            print('creating output directory...')
            os.mkdir(output_dir)
        os.chdir(output_dir)
        for_csv = self.to_pandas()
        for_csv.to_csv(path_or_buf = output_dir + '/' +
                       self.filename[0:self.filename.index('.tif')]+ '.csv',
                       index = True, header = True)
    def rm_border_objs(self, border = 1, z = True):
        '''remove all objects that contact the edge of the 3D stack.

        Args:
            border (int, optional): the size of the border around the edge which a pixel from
               the object must contact to be removed. Defaults to 1.
            z (bool, optional): should objects that contact the z-axis edge be eliminated? if
               True, any object with a pixel in the top or bottom image of the
               stack is removed. Defaults to True.

        Output: alters the objects within the MitoSegmentObj. Removes objects
        from the mitochondria image, the obj_nums, and all other variables with
        elements of obj_nums as keys in a dict (parents, volumes, etc)
        '''

        border_mask = np.full(shape = self.mitochondria.shape, fill_value = True,
                              dtype = bool)
        if z == True:
            border_mask[border:-border,border:-border,border:-border] = False
        elif z == False:
            border_mask[:,border:-border,border:-border] = False
        objs_to_rm = np.unique(self.mitochondria[border_mask])
        objs_to_rm = objs_to_rm[objs_to_rm != 0]
        for x in objs_to_rm:
            self.mitochondria[self.mitochondria == x] = 0
            self.obj_nums.remove(x)
            self.volumes.pop(x, None)
            if hasattr(self, "parent"):
                self.parent.pop(x, None)
        self.nmitos = len(self.obj_nums)
        self.border_rm_flag = True

    ## HELPER METHODS ##

    def to_pandas(self):
        '''create a pandas DataFrame of tabulated numeric data.
        
        the pdout attribute indicates which variables to include in the
        DataFrame.
        
        Helper method for self.to_csv().
        '''
        df_dict = {}
        for attr in self.pdout:
            df_dict[str(attr)] = pd.Series(getattr(self, attr))
        if 'volumes' in self.pdout:
            vflag_out = dict(zip(self.obj_nums,
                                 [self.volumes_flag]*len(self.obj_nums)))
            df_dict['volumes_flag'] = pd.Series(vflag_out)
        return pd.DataFrame(df_dict)
    def convert_volumes(self, z = 0.2, x = 0.0675):
        '''convert volumes from units of pixels to metric units.

        The default values provided correspond to the appropriate values for
        imaging using the Murray spinning disk confocal microscope with 0.2 um
        spacing in between slices.

        Args:
            z (float, optional): the distance between slices in the z-stack in
                units of microns.
            x (float, optional): the linear distance between adjacent pixels in
                each slice in units of microns. x is also used for y.

        Output: converts self.volumes to units of femtoliters, and changes the
            self.volumes_flag to 'femtoliters'.
        '''
        conv_factor = z*x*x
        for key, val in self.volumes:
            self.volumes[key] = float(self.volumes[key])*conv_factor
        self.volumes_flag = 'femtoliters' 
    def plot_stack(self, stack_arr, colormap='jet'):
        ''' Create a matplotlib plot with each subplot containing a slice.
        
        Args:
            stack_arr (ndarray of ints): A numpy ndarray containing pixel 
                intensity values.
            colormap: The colormap to be used when displaying pixel intensities.
                Defaults to jet.
        
        Yield: a pyplot object in which each slice from the image array
            is represented in a subplot. subplots are 4 columns
            across (when 4 or more slices are present) with rows to
            accommodate all slices.
        '''

        nimgs = stack_arr.shape[0] # z axis of array dictates number of slices
        # plot with 4 imgs across

        # determine how many rows and columns of images there are

        if nimgs < 5:
            f, axarr = plt.subplots(1,nimgs)

            for i in range(0,nimgs):
                axarr[i].imshow(stack_arr[i,:,:], cmap=colormap)
                axarr[i].xaxis.set_visible(False)
                axarr[i].yaxis.set_visible(False)
            f.set_figwidth(16)
            f.set_figheight(4)        

        else:
            f, axarr = plt.subplots(int(np.ceil(nimgs/4)),4)

            for i in range(0,nimgs):
                r = int(np.floor(i/4))
                c = int(i % 4)
                axarr[r,c].imshow(stack_arr[i,:,:], cmap=colormap)
                axarr[r,c].xaxis.set_visible(False)
                axarr[r,c].yaxis.set_visible(False)
        # manage remainder images
            if nimgs%4 > 0:
                r = int(np.floor(nimgs/4))

                for c in range(nimgs%4,4):
                    axarr[r,c].axis('off')

            f.set_figwidth(16)
            f.set_figheight(4*np.ceil(nimgs/4))

    def plot_maxima_stack(self, masked_max, smooth_dist):

        ''' Creates a matplotlib plot object in which each slice from the image
        is displayed as a single subplot, in a 4-by-n matrix (n depends upon
        the number of slices in the image)'''

        nimgs = masked_max.shape[0] # z axis of array dictates number of slices
        # plot with 4 imgs across

        # determine how many rows and columns of images there are

        if nimgs < 5:
            f, axarr = plt.subplots(1,nimgs)

            for i in range(0,nimgs):
                axarr[i].imshow(smooth_dist[i,:,:], cmap='gray')
                axarr[i].imshow(masked_max[i,:,:], cmap='autumn')
                axarr[i].xaxis.set_visible(False)
                axarr[i].yaxis.set_visible(False)
            f.set_figwidth(16)
            f.set_figheight(4)

        else:
            f, axarr = plt.subplots(int(np.ceil(nimgs/4)),4)

            for i in range(0,nimgs):
                r = int(np.floor(i/4))
                c = int(i%4)
                axarr[r,c].imshow(smooth_dist[i,:,:], cmap='gray')
                axarr[r,c].imshow(masked_max[i,:,:], cmap='autumn')
                axarr[r,c].xaxis.set_visible(False)
                axarr[r,c].yaxis.set_visible(False)

            if nimgs%4 > 0:
                r = int(np.floor(nimgs/4))

                for c in range(nimgs%4, 4):
                    axarr[r,c].axis('off')

            f.set_figwidth(16)
            f.set_figheight(4*np.ceil(nimgs/4))


class MitoSegmenter:
 
    '''Class for segmenting reticular shapes in multipage TIFF images.
    
    When called, generates an object of class MitoSegmenter with the
    segmentation parameters defined as per Args for the image indicated. The
    segment() method can then be called to perform segmentation, returning a
    MitoSegmentObj.

    Args:
        filename (str): The filename of the multipage TIFF-format image to be
        segmented. This is defined with respect to the current working 
        directory.
        seg_method (str, optional): The segmentation approach to use in a string
        format. This can be one of the following:
            threshold (default): Segmentation by setting an absolute cutoff for
                minimum pixel intensity of an object. Requires kwargs dependent
                upon the segmentation mode (see mode below).
            canny: Segmentation by Canny edge detection. Requires kwargs
                high_threshold and low_threshold to be set, or the default
                values 1000 and 500 will be used.
        mode (str, optional): The segmentation mode to be used if seg_method ==
            threshold. This can be one of the following:
            threshold (default): Uses a user-provided pixel intensity value as
                the segmentation threshold. Requires kwarg threshold.
            bg_scaled: Use the median value from all pixels that correspond to
                cells within a CellSegmentObj (see CellSegment.py in this
                module) segmented from another fluorescence channel as the
                background. Requires kwargs cells and bg_diff.

    Additional kwargs:
        threshold (int): Only relevant if using seg_method threshold and mode
            threshold. The pixel intensity value to be used as a cutoff for
            segmented objects.
        cells (str): Only relevant if using seg_method threshold and mode
            bg_scaled. The name of a CellSegmentObj already present in the
            environment which corresponds to cells segmented from the same
            field and in a separate channel. See CellSegment.py for details.
        bg_diff (int): Only relevant if using seg_method threshold and mode
            bg_scaled. The desired pixel intensity units above the cells' 
            median pixel intensity to set the threshold.
        high_threshold (int, optional): Only relevant if using seg_method 
            "canny". The high threshold to be passed to skimage's canny edge
            detector. Default value is 1000. See skimage.feature.canny 
            documentation for more details.
        low_threshold (int, optional): Only relevant if using seg_method 
            "canny". The low threshold to be passed to skimage's canny edge
            detector. Default value is 500. See skimage.feature.canny 
            documentation for more details.
        min_cutoff (int, optional): Only relevant if using seg_method 'canny'.
            An integer value for the minimal brightness that can be assigned to
            a mitochondrion object, which allows for re-opening holes in
            donut-shaped mitochondria. See MitoSegmenter.segment() code for
            details.

    Attributes:
        filename (str): The filename provided in args.
        seg_method (str): The segmentation method provided in args, defaults to
            'threshold'.
        mode (str): The segmentation mode provided in args, defaults to
            'threshold'.
        Additional attributes provided by **kwargs (see Args, Additional kwargs
            above).

    '''
   
    def __init__(self,filename, seg_method = 'threshold', mode = 'threshold',
                 **kwargs):
        self.filename = filename
        self.seg_method = seg_method
        self.mode = mode
        for key, value in kwargs.items():
            setattr(self,key,value)
        if self.seg_method == 'canny':
            self.high_threshold = kwargs.get('high_threshold',1000)
            self.low_threshold = kwargs.get('low_threshold',500)
        if self.seg_method == 'threshold':
            if mode == 'threshold':
                self.threshold = kwargs.get('threshold',float('nan'))
                if np.isnan(self.threshold):
                    raise ValueError('A threshold argument must be provided to segment with a constant threshold.')
            if mode == 'bg_scaled':
                self.cells = kwargs.get('cells', '')
                self.bg_diff = float(kwargs.get('bg_diff',float('nan')))
                if self.cells == '':
                    raise ValueError('A CellSegmentObj containing segmented cells is required if mode == bg_scaled.')
                if np.isnan(self.bg_diff):
                    raise ValueError('a bg_diff argument is needed if mode == bg_scaled.')
    def segment(self):
        '''Segment objects within the image according to attributes provided.
        Yields: a MitoSegmentObj containing segmented objects as well as all
            images generated during segmentation (for post-hoc analysis) as
            well as relevant values, e.g. numbers and names of segmented
            particles. See MitoSegmentObj documentation for more details.
        '''        
        starttime = time.time() # begin timing
        f_directory = os.getcwd()
        pdout = [] # list of MitoSegmentObj attributes to pass to pandas for csv
        # data import
        print('reading' + self.filename)
        raw_img = io.imread(self.filename)
        print('raw image imported.')
        # gaussian filter assuming 100x objective and 0.2 um slices
        print('performing gaussian filtering...')
        # use empirically determined optimal values for the gaussian filter
        gaussian_img = gaussian_filter(raw_img, [0.75,0.75,0.75])
        print('cytosolic image smoothed.')
        print('preprocessing complete.')
        ## SEGMENTATION BY THRESHOLDING THE GAUSSIAN ##
        if self.seg_method == 'threshold':
            # binary thresholding and cleanup
            print('thresholding...')
            threshold_img = np.copy(gaussian_img)
            if self.mode == 'threshold':
               print('mode = threshold.')
               # make binary image
               threshold_img[threshold_img < self.threshold] = 0
               threshold_img[threshold_img > 0] = 1
               print('thresholding complete.')
            if self.mode == 'bg_scaled':
                print('mode = background-scaled.')
                self.thresholds = {}
                threshold_img = np.zeros(shape = raw_img.shape)
                for i in self.cells.obj_nums:
                    if i == 0:
                        pass
                    else:
                        print('thresholding cell ' + str(i))
                        # get median for the cell
                        cell_median = np.median(gaussian_img[self.cells.final_cells == i])
                        # generate the thresholded binary mask for each cell
                        threshold_img[np.logical_and(self.cells.final_cells == i,
                                      gaussian_img > cell_median + self.bg_diff)] = 1
                        self.thresholds[i] = cell_median + self.bg_diff # store val
                print('thresholding complete.')
            # distance and maxima transformation to find objects
            # next two steps assume 100x objective and 0.2 um slices
            print('generating distance map...')
            dist_map = distance_transform_edt(threshold_img, sampling = (2,1,1))
            print('distance map complete.')
            print('smoothing distance map...')
            smooth_dist = gaussian_filter(dist_map, [1,2,2])
            print('distance map smoothed.')
            print('identifying maxima...')
            # find local maxima in the smoothed distance map
            # these will be the watershed seeds
            max_strel = generate_binary_structure(3,2)
            maxima = maximum_filter(smooth_dist,
                                    footprint = max_strel) == smooth_dist
            # clean up background and edges
            bgrd_3d = smooth_dist == 0
            eroded_bgrd = binary_erosion(bgrd_3d, structure = max_strel,
                                         border_value = 1)
            maxima = np.logical_xor(maxima, eroded_bgrd)
            print('maxima identified.')
            # watershed segmentation
            labs = self.watershed_labels(maxima)
            print('watershedding...')
            mitochondria = watershed(-smooth_dist, labs, mask = threshold_img)
            print('watershedding complete.')
            if self.mode == 'bg_scaled':
                # find cell boundaries and define objects that are on the
                # edges, then assign segmented objects to parent cells
                edge_struct = generate_binary_structure(3,1)
                self.c_edges = {}
                print('finding edges of cells...')
                for i in self.cells.obj_nums:
                    self.c_edges[i] = np.logical_xor(self.cells.final_cells == i,
                                                          binary_erosion(self.cells.final_cells== i,
                                                                         edge_struct))
                print('cell edges found.')
                self.primary_objs = [x for x in np.unique(mitochondria) if x != 0]
                self.parent = {}
                self.obj_edges = {}
                self.on_edge = {}
                mito_mask = mitochondria != 0
                for obj in self.primary_objs:
                    self.parent[obj] = self.cells.final_cells[labs == obj][0]
                    obj_mask = mitochondria == obj
                    obj_edge = np.logical_xor(obj_mask, 
                                              binary_erosion(obj_mask,
                                                             edge_struct))
                    self.obj_edges[obj] = obj_edge
                    # test if the object's edge and its cell's edge overlap
                    if np.any(np.logical_and(obj_edge,
                                             self.c_edges[self.parent[obj]])):
                        self.on_edge[obj] = True
                        print('object on the edge: ' + str(obj))
                        print('parent cell: ' + str(self.parent[obj]))
                        new_obj = obj_mask
                        search_obj = obj_mask
                        tester = 0
                        iteration = 1
                        while tester == 0:
                            # TODO: FIX THIS BLOCK OF CODE! GETTING STUCK WITHIN
                            # IT! NOT SURE HOW MANY ITERATIONS ITS DOING, OR FOR
                            # HOW MANY DIFFERENT PEROXISOMES.
                            new_px = binary_dilation(search_obj, edge_struct)
                            new_px[np.logical_or(new_obj, mito_mask)] = False
                            print('iteration: ' + str(iteration))
                            # print('new pixels for iteration ' + str(iteration) + \
                            #      ': ')
                            # print(np.nonzero(new_px))
                            if np.any(gaussian_img[new_px] >
                                      self.thresholds[self.parent[obj]]):
                                to_add = np.logical_and(new_px, gaussian_img >
                                                        self.thresholds[self.parent[obj]])
                                new_obj = np.logical_or(new_obj, to_add)
                            #    print('object pixels after iteration '
                            #          + str(iteration) + ': ')
                            #    print(np.nonzero(new_obj))
                                search_obj = to_add # only search from new pixels
                            else:
                                mitochondria[new_obj] = obj
                                tester = 1
                            iteration = iteration + 1
                    else:
                        self.on_edge[obj] = False
        elif self.seg_method == 'canny':
            print('performing Canny edge detection-based segmentation')
            ## EDGE-DETECTION BASED SEGMENTATION ##
            threshold_img = np.empty_like(gaussian_img)
            edge_img = np.empty_like(gaussian_img)
            c_strel = generate_binary_structure(2,1)
            # perform canny edge detection on each slice s
            for s in range(0,gaussian_img.shape[0]):
                print('performing Canny edge detection on slice ' + str(s))
                c = canny(gaussian_img[s,:,:],
                          sigma = 0,
                          low_threshold = self.low_threshold,
                          high_threshold = self.high_threshold)
                # clean up objects edges that have gaps
                print('cleaning up slice ' + str(s))
                c = binary_closing(c,c_strel)
                edge_img[s,:,:] = np.copy(c)
                # fill holes to generate binary masks of objects
                c = binary_fill_holes(c)
                c = binary_opening(c, c_strel) # eliminate incomplete lines
                c = binary_closing(c, c_strel) # clean up edges
                threshold_img[s,:,:] = c
            if self.min_cutoff:
                # re-open holes in donut-shaped objects by eliminating pixels
                # with values < min_cutoff
                print('generating holes using a minimum intensity cutoff')
                print('of ' + str(self.min_cutoff))
                threshold_img[gaussian_img < self.min_cutoff] = 0
            # Occasionally, this code failed to find edges in individual
            # slices, but found objects in the slice immediately above and
            # below, due to edge strength issues. To remedy this, gaps were
            # filled when an object was found in the slice below and above at a
            # given pixel.
            v_merge_strel = np.array([[[0,0,0],
                                       [0,1,0],
                                       [0,0,0]],
                                      [[0,0,0],
                                       [0,1,0],
                                       [0,0,0]],
                                      [[0,0,0],
                                       [0,1,0],
                                       [0,0,0]]])
            print('closing holes between slices...')
            threshold_img = binary_closing(threshold_img,
                                           structure=v_merge_strel)
            print('generating distance map...')
            dist_map = distance_transform_edt(threshold_img, sampling = (3,1,1))
            print('distance map complete.')
            print('smoothing distance map...')
            smooth_dist = gaussian_filter(dist_map, [1,2,2])
            print('distance map smoothed.')
            print('identifying maxima...')
            max_strel = generate_binary_structure(3,2)
            # identify local maxima (these will be the seed points for
            # watershed segmentation)
            maxima = maximum_filter(smooth_dist,
                                    footprint = max_strel) == smooth_dist
            # clean up background and edges
            bgrd_3d = smooth_dist == 0
            eroded_bgrd = binary_erosion(bgrd_3d, structure = max_strel,
                                         border_value = 1)
            maxima = np.logical_xor(maxima, eroded_bgrd)
            print('maxima identified.')
            # watershed segmentation
            labs = self.watershed_labels(maxima)
            print('watershedding...')
            raw_wsheds = watershed(-smooth_dist, labs, mask = threshold_img)
            print('watershedding complete.')
            print('merging adjacent watershed objects...')
            # merge contiguous objects into one longer reticular structure
            mitochondria = self.merge_wsheds(raw_wsheds)
            print('watershed merging complete.')
            if hasattr(self,'cells'):
                # assign segmented objects to cells if a CellSegmentObj was
                # included
                self.primary_objs = [x for x in np.unique(mitochondria) \
                                     if x != 0]
                self.parent = {}
                for obj in self.primary_objs:
                    o_parent = self.cells.final_cells[labs == obj][0]
                    if o_parent == 0:
                        self.primary_objs.remove(obj)
                    else:
                        self.parent[obj] = o_parent
        obj_nums, volumes = np.unique(mitochondria, return_counts = True)
        volumes=dict(zip(obj_nums.astype('uint16'),volumes))
        # remove background from volumes and obj_nums
        del volumes[0]
        obj_nums = obj_nums.astype('uint16').tolist()
        obj_nums.remove(0)
        # generate dict of relevant parameters to pass to MitoSegmentObj
        mode_params = {}
        if hasattr(self, 'parent'):
            pdout.append('parent')
            mode_params['parent'] = self.parent
        if self.seg_method == 'canny':
            mode_params['high_threshold'] = self.high_threshold
            mode_params['low_threshold'] = self.low_threshold
            mode_params['edges'] = edge_img
            pdout.append('volumes')
        if self.seg_method == 'threshold':
            if self.mode == 'threshold':
                mode_params['threshold'] = self.threshold
                pdout.append('volumes')
            elif self.mode == 'bg_scaled':
                mode_params['thresholds'] = self.thresholds
                mode_params['bg_diff'] = self.bg_diff
                mode_params['cells'] = self.cells
                mode_params['cell_edges'] = self.c_edges
                mode_params['cell_nums'] = self.cells.obj_nums
                mode_params['obj_edges'] = self.obj_edges
                mode_params['on_edge'] = self.on_edge
                for x in ['thresholds','on_edge','parent', 'volumes']:
                    pdout.append(x)
        return MitoSegmentObj(f_directory, self.filename, raw_img,
                             gaussian_img, self.seg_method, self.mode, 
                             threshold_img, dist_map, smooth_dist, maxima,
                             labs, mitochondria, obj_nums, volumes,
                             to_pdout = pdout, mode_params = mode_params)

    ## HELPER METHODS ##
    def watershed_labels(self, maxima_img):
        '''Number local maxima in order for use in watershedding.

        Args:
            maxima_img (np.ndarray): A boolean array with local maxima labeled 
            as true pixels.
        
        Yields:
            A numpy ndarray with maxima numbered sequentially.
        '''
        max_z, max_y, max_x = np.nonzero(maxima_img)
        label_output = np.zeros(maxima_img.shape)
        for i in range(0,len(max_y)):
            label_output[max_z[i],max_y[i],max_x[i]] = i+1
        return(label_output) 
    def merge_wsheds(self, wshed_img):
        '''Merge adjacent watersheds in raw watershed output.
        
        Because of the reticular shape of mitochondria, many local maxima
        within a single contiguous mitochondria are often identified by the
        Euclidean distance transform, resulting in many watershedded objects
        within one contiguous mitochondrion. This method iteratively merges
        contiguous objects until there are no contiguous objects segmented into
        separate watersheds.
        
        Args:
            wshed_img (nd_array of ints): The raw watershedding output from
                MitoSegmenter.segment().
        
        Yields:
            temp_img (ndarray): The same image with all contiguous objects
                merged to share one object ID number.'''
        obj_nums = np.unique(wshed_img)
        temp_img = np.copy(wshed_img) # generate a copy to modify
        # generate a structuring element to expand objects
        dil_strel = np.array([[[0,0,0],
                               [0,1,0],
                               [0,0,0]],
                              [[0,0,0],
                               [0,1,0],
                               [0,0,0]],
                              [[0,1,0],
                               [1,1,1],
                               [0,1,0]],
                              [[0,0,0],
                               [0,1,0],
                               [0,0,0]],
                              [[0,0,0],
                               [0,1,0],
                               [0,0,0]]]) # increase connectivity across Z axis
        for n in obj_nums:
            # skip background
            if n == 0:
               continue 
            else:
                # test if the object number has already been removed by merging
                # with another watershed earlier
                if not np.any(temp_img == n):
                    continue
                tester = False # used to break out of the next while loop
                iteration = 1 # for counting/recording purposes
                print('testing object #' + str(n))
                while not tester:
                    print('iteration #' + str(iteration))
                    # generate a bool array where the object is
                    n_img = temp_img == n
                    dil_n = binary_dilation(n_img) #expand to contacting pixels
                    overlap_objs = np.unique(temp_img[dil_n]) # find objs that overlap
                    print('overlapping objects:')
                    print(overlap_objs)
                    overlap_objs = overlap_objs[overlap_objs != 0] #rm bgrd px
                    # if objects other than n are present in this set
                    if overlap_objs.size > 1:
                        # replace all values for the overlapping objects with n
                        # (merge the objects)
                        temp_img[np.in1d(temp_img,overlap_objs).reshape(temp_img.shape)] = n
                        iteration = iteration + 1
                    else:
                        # break out of the while loop
                        tester = True
        return temp_img
                        
