
"""
This python script is a wrapper for Simple ITK's watershed
Copyright (C) 2020 SPAM Contributors

This program is free software: you can redistribute it and/or modify it
under the terms of the GNU General Public License as published by the Free
Software Foundation, either version 3 of the License, or (at your option)
any later version.

This program is distributed in the hope that it will be useful, but WITHOUT
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
more details.

You should have received a copy of the GNU General Public License along with
this program.  If not, see <http://www.gnu.org/licenses/>.
"""

import spam.label
import spam.helpers
import spam.mesh

import tifffile
import numpy
numpy.seterr(all='ignore')
import scipy.ndimage
import argparse

# Define argument parser object
parser = argparse.ArgumentParser()

# Parse arguments with external helper function
args = spam.helpers.optionsParser.ITKwatershedParser(parser)
if args.VERBOSE:
    print("-> Loading binary image...", end='')
binVol = tifffile.imread(args.inFile.name)
if args.VERBOSE:
    print("done.")

if args.MARKER_FILE is not None:
    if args.VERBOSE:
        print("-> Loading marker image...", end='')
    markerVol = tifffile.imread(args.MARKER_FILE)
    if args.VERBOSE:
        print("done.")

    # 2019-09-07 EA: changing dilation/erosion into a single pass by a spherical element, rather than repeated
    # iterations of the standard.
    if args.LABEL_DILATE != 0:
        BB = spam.label.boundingBoxes( markerVol)
        COM = spam.label.centresOfMass(markerVol, boundingBoxes=BB)
        #plt.imshow(imLab[25]); plt.show()
        tmp = numpy.zeros_like(markerVol)
        for label in range(1,markerVol.max()+1):
            gl = spam.label.getLabel(markerVol, label, labelDilate=args.LABEL_DILATE, boundingBoxes=BB, centresOfMass=COM)
            if gl is not None:
                tmp[gl['slice']] = gl['subvol']*label
    markerVol = tmp
else:
    markerVol = None

lab = spam.label.watershed(binVol, markers=markerVol)
if args.VERBOSE:
    print("-> Saving labelled image...", end='')
tifffile.imsave(args.OUT_DIR+"/"+args.PREFIX+".tif", lab)
if args.VERBOSE:
    print("done.")
