#! /usr/bin/env python
# -*- coding: utf-8 -*
##########################################################################
# NSAp - Copyright (C) CEA, 2017
# Distributed under the terms of the CeCILL-B license, as published by
# the CEA-CNRS-INRIA. Refer to the LICENSE file or to
# http://www.cecill.info/licences/Licence_CeCILL-B_V1-en.html
# for details.
##########################################################################

# System import
import os
import json
import argparse
from datetime import datetime
from pprint import pprint
import textwrap
from argparse import RawTextHelpFormatter

# Bredala module
try:
    import bredala
    bredala.USE_PROFILER = False
    bredala.register("pydcmio.deface.mask_face", names=["deface"])
    bredala.register("pydcmio.deface.mri_deface", names=["deface"])
    bredala.register("pydcmio.deface.pdeface", names=["deface"])
except:
    pass

# Package import
from pydcmio import __version__ as version
from pydcmio.deface import mask_face
from pydcmio.deface import mri_deface
from pydcmio.deface import pdeface

# Pyconnectome import
from pyconnectome import DEFAULT_FSL_PATH
from pyconnectome import __version__ as pyconnectome_version
from pyconnectome.wrapper import FSLWrapper

# Pyfreesurfer import
from pyfreesurfer import DEFAULT_FREESURFER_PATH
from pyfreesurfer import __version__ as pyfreesurfer_version
from pyfreesurfer.wrapper import FSWrapper


# Parameters to keep trace
__hopla__ = ["runtime", "inputs", "outputs"]


# Script documentation
DOC = """
Defacing with the mask_face command
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Wraps around the 'mask_face' command.

This code enables us to deface Nifti images using the Marcus, D.
'mask_face' command.

The code is setup so that the mask_face matlab dependencies
have been compiled. Thus you need the MathWorks MCR software installed
on your computer to run this script.

Steps:

1- Convert input Nifti images to .img/.hdr Analyze images.
2- Start defacing with the 'mask_face' command.

MaskFace command example:

python $HOME/git/pydcmio/pydcmio/scripts/pydcmio_maskface \
    -i $HOME/Desktop/results/sample_T1_input.nii.gz \
    -o $HOME/Desktop/results/001 \
    -T mask_face \
    -F /etc/fsl/5.0/fsl.sh \
    -M /neurospin/nsap/local/maskface/MCR/v901 \
    -v 2 \
    -E

FreeSurfer 'mri_deface' command example:

python $HOME/git/pydcmio/pydcmio/scripts/pydcmio_maskface \
    -i $HOME/Desktop/results/sample_T1_input.mgz \
    -o $HOME/Desktop/results/002 \
    -T mri_deface \
    -H /i2bm/local/freesurfer/SetUpFreeSurfer.sh \
    -v 2

PyDeface command example:

python $HOME/git/pydcmio/pydcmio/scripts/pydcmio_maskface \
    -i $HOME/Desktop/results/sample_T1_input.nii.gz \
    -o $HOME/Desktop/results/003 \
    -T pydeface \
    -F /etc/fsl/5.0/fsl.sh \
    -v 2 
"""


def is_file(filearg):
    """ Type for argparse - checks that file exists but does not open.
    """
    if not os.path.isfile(filearg):
        raise argparse.ArgumentError(
            "The file '{0}' does not exist!".format(filearg))
    return filearg


def is_directory(dirarg):
    """ Type for argparse - checks that directory exists.
    """
    if not os.path.isdir(dirarg):
        raise argparse.ArgumentError(
            "The directory '{0}' does not exist!".format(dirarg))
    return dirarg


def get_cmd_line_args():
    """
    Create a command line argument parser and return a dict mapping
    <argument name> -> <argument value>.
    """
    parser = argparse.ArgumentParser(
        prog="python pydcmio_maskface",
        description=textwrap.dedent(DOC),
        formatter_class=RawTextHelpFormatter)

    # Required arguments
    required = parser.add_argument_group("required arguments")
    required.add_argument(
        "-i", "--images",
        required=True, nargs="+",
        help="Input MRI head images to be defaced.")
    required.add_argument(
        "-o", "--outdir",
        required=True, type=is_directory, metavar="<path>",
        help="The destination folder.")

    # Optional arguments
    parser.add_argument(
        "-R", "--reference",
        type=is_file,
        help="The image that must be used as reference if more than one image"
             "have been supplied as input or if the input volume is 4d.")
    parser.add_argument(
        "-T", "--type",
        dest="defacing_type", choices=("mask_face", "mri_deface", "pydeface"),
        help="The defacing method to be applied.")
    parser.add_argument(
        "-E", "--ears",
        action="store_true",
        help="If activated, masks ears (only valid for 'mask_face').")
    parser.add_argument(
        "-H", "--fs-sh",
        type=is_file, metavar="<path>",
        help="Bash script initializing FreeSurfer's environment.")
    parser.add_argument(
        "-F", "--fsl-sh",
        type=is_file, metavar="<path>",
        help="Bash script initializing FSL's environment.")
    parser.add_argument(
        "-M", "--matlabmcr",
        type=is_directory,
        help="The MATLAB MCR directory for standalone applications (only "
             "required for 'mask_face').")
    parser.add_argument(
        "-K", "--keep",
        dest="keep_workspace", action="store_true",
        help="If activated, keep the defacing workspace (require more "
             "disk space).")
    parser.add_argument(
        "-v", "--verbose",
        type=int, choices=[0, 1, 2],
        help="Increase the verbosity level: 0 silent, [1, 2] verbose.")

    # Create a dict of arguments to pass to the 'main' function
    args = parser.parse_args()
    kwargs = vars(args)
    verbose = kwargs.pop("verbose")
    if kwargs["fs_sh"] is None:
        kwargs["fs_sh"] = DEFAULT_FREESURFER_PATH
    if kwargs["fsl_sh"] is None:
        kwargs["fsl_sh"] = DEFAULT_FSL_PATH

    return kwargs, verbose


"""
Parse the command line.
"""
inputs, verbose = get_cmd_line_args()
tool = "pydcmio_maskface"
timestamp = datetime.now().isoformat()
tool_version = version
fsl_version, freesurfer_version = (None, None)
if inputs["defacing_type"] == "mask_face":
    fsl_version = FSLWrapper([], shfile=inputs["fsl_sh"]).version
elif inputs["defacing_type"] == "mri_deface":
    freesurfer_version = FSWrapper([], inputs["fs_sh"]).version
params = locals()
runtime = dict([(name, params[name])
               for name in ("tool","tool_version", "fsl_version",
                            "freesurfer_version", "timestamp", 
                            "pyconnectome_version", "pyfreesurfer_version")])
outputs = None
if verbose > 0:
    print("[info] Starting defacing ...")
    print("[info] Runtime:")
    pprint(runtime)
    print("[info] Inputs:")
    pprint(inputs)


"""
Defacing
"""
if inputs["defacing_type"] == "mask_face":
    deface_files, snap_files = mask_face(
        input_files=inputs["images"],
        outdir=inputs["outdir"],
        reference_file=inputs["reference"],
        mask_ears=inputs["ears"],
        fsl_sh=inputs["fsl_sh"],
        matlab_mcr=inputs["matlabmcr"],
        verbose=verbose,
        rm_workspace=(not inputs["keep_workspace"]))
elif inputs["defacing_type"] == "mri_deface":
    deface_files, snap_files = mri_deface(
        input_files=inputs["images"],
        outdir=inputs["outdir"],
        reference_file=None,
        verbose=verbose,
        fs_config=inputs["fs_sh"])
elif inputs["defacing_type"] == "pydeface":
    deface_files, snap_files = pdeface(
        input_files=inputs["images"],
        outdir=inputs["outdir"],
        reference_file=None,
        fsl_sh=inputs["fsl_sh"])
else:
    raise ValueError("Invalid defacing type.")


"""
Update the outputs and save them and the inputs in a 'logs' directory.
"""
logdir = os.path.join(inputs["outdir"], "logs")
if not os.path.isdir(logdir):
    os.mkdir(logdir)
params = locals()
outputs = dict([(name, params[name])
               for name in ("deface_files", "snap_files")])
for name, final_struct in [("inputs", inputs), ("outputs", outputs),
                           ("runtime", runtime)]:
    log_file = os.path.join(logdir, "{0}.json".format(name))
    with open(log_file, "wt") as open_file:
        json.dump(final_struct, open_file, sort_keys=True, check_circular=True,
                  indent=4)
if verbose > 1:
    print("[info] Outputs:")
    pprint(outputs)
