#!python

# 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 numpy
import spam.helpers
import spam.label
import spam.mesh
import tifffile

numpy.seterr(all="ignore")
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.imwrite(args.OUT_DIR + "/" + args.PREFIX + ".tif", lab)
if args.VERBOSE:
    print("done.")
