#!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 spam.helpers
import spam.label as ltk
import spam.DIC
import spam.deformation
import os

os.environ["OPENBLAS_NUM_THREADS"] = "1"
import numpy

numpy.seterr(all="ignore")
import sys, os
import tifffile

# import matplotlib.pyplot as plt
# Added from rotation_correlation-style parallelisation
import multiprocessing
try:                 multiprocessing.set_start_method('fork')
except RuntimeError: pass
import argparse
import progressbar


# Define argument parser object
parser = argparse.ArgumentParser(
    description="spam-moveLabels "
    + spam.helpers.optionsParser.GLPv3descriptionHeader
    + "This script applies discretely measured deformation functions (Phi) coming from 'spam-ddic' to a labelled image, "
    + "thus generating the deformed labelled image.\n\nWarning: since we're moving labels, "
    + "nearest neighbour interpolation must be used, and thus the shapes of the labels will be slightly damaged",
    formatter_class=argparse.RawTextHelpFormatter,
)

# Parse arguments with external helper function
args = spam.helpers.optionsParser.moveLabelsParser(parser)

if args.PROCESSES is None:
    args.PROCESSES = multiprocessing.cpu_count()

print("spam-moveLabels -- Current Settings:")
argsDict = vars(args)
for key in sorted(argsDict):
    print("\t{}: {}".format(key, argsDict[key]))

if args.RETURN_STATUS_THRESHOLD is None:
    DVC = spam.helpers.readCorrelationTSV(args.TSVFile.name, readConvergence=False)
    RS = None
else:
    DVC = spam.helpers.readCorrelationTSV(args.TSVFile.name, readConvergence=True)
    RS = DVC["returnStatus"]

# Read labelled image
lab = tifffile.imread(args.LabFile.name)


# if args.GREY_FILE is not None:
# greyTmp = tifffile.imread(args.GREY_FILE)
# grey = numpy.zeros((labShapeOrig[0]+2*args.OVERALL_MARGIN,
# labShapeOrig[1]+2*args.OVERALL_MARGIN,
# labShapeOrig[2]+2*args.OVERALL_MARGIN), dtype=greyTmp.dtype)
# grey[slicePadToNonPad] = greyTmp

labOut = spam.label.moveLabels(
    lab,
    DVC["PhiField"],
    returnStatus=RS,
    margin=args.MARGIN,
    PhiCOM=args.PHICOM,
    threshold=args.THRESH,
    labelDilate=args.LABEL_DILATE,
    nProcesses=args.PROCESSES,
)

# if args.GREY_FILE is None:
print("\nSaving labelled image with displaced grains...", end="")
tifffile.imwrite(args.OUT_DIR + "/" + args.PREFIX + ".tif", labOut.astype(lab.dtype))
# else:
# print("\nSaving binary image with displaced grains...", end='')
# tifffile.imwrite( args.OUT_DIR+"/"+args.PREFIX+".tif", labOut[slicePadToNonPad].astype('<u1')*255 )
print("done.")
