#! /usr/bin/env python
##########################################################################
# NSAp - Copyright (C) CEA, 2013 - 2016
# 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
from __future__ import print_function
import argparse
import os
import shutil

# Bredala import
try:
    import bredala
    bredala.USE_PROFILER = False
    bredala.register("pyconnectomist.preproc",
                     names=["complete_preprocessing"])
    bredala.register("pyconnectomist.preproc.qsapce",
                     names=["data_import_and_qspace_sampling"])
    bredala.register("pyconnectomist.preproc.mask",
                     names=["rough_mask_extraction"])
    bredala.register("pyconnectomist.preproc.outliers",
                     names=["outlying_slice_detection"])
    bredala.register("pyconnectomist.preproc.susceptibility",
                     names=["susceptibility_correction"])
    bredala.register("pyconnectomist.preproc.registration",
                     names=["dwi_to_anatomy"])
    bredala.register("pyconnectomist.preproc.eddy",
                     names=["eddy_and_motion_correction",
                            "export_eddy_motion_results_to_nifti"])
    bredala.register("pyconnectomist.wrappers",
                     names=["ConnectomistWrapper.__call__"])
except:
    pass

# Pyconnectomist import
from pyconnectomist import __version__ as version
from pyconnectomist.preproc import complete_preprocessing
from pyconnectomist.manufacturers import MANUFACTURERS
from pyconnectomist.wrappers import ConnectomistWrapper
from pyconnectomist import DEFAULT_CONNECTOMIST_PATH


# Parameters to keep trace
__hopla__ = ["tool", "tool_version", "version", "config", "inputs", "outputs",
             "outdir", "subjectid", "subjdir", "dwi", "bvec", "bval",
             "manufacturer", "delta_te", "partial_fourier_factor",
             "parallel_acceleration_factor", "b0_magnitude", "b0_phase",
             "b0_magnitude", "invertx", "inverty", "invertz", "negative_sign",
             "echo_spacing", "epi_factor", "b0_field", "water_fat_shift",
             "delete_steps", "morphologist_dir"]

# Script documentation
doc = """
Connectomist preproc
~~~~~~~~~~~~~~~~~~~~

Function that runs all the Connectomist preprocessing tabs.
Generates results in '<outdir>/<subjectid>/proeproc'.

Steps:

1- Create the preprocessing output directory if not existing.
2- Import files to Connectomist and choose q-space model.
3- Create a brain mask.
4- Detect and correct outlying diffusion slices.
5- Susceptibility correction.
6- Eddy current and motion correction.
7- Export result as a Nifti with a .bval and a .bvec.
8- Export outliers.
9- Registration t1 - dwi
10- Delete intermediate files and directories if requested.

Command:

python $HOME/git/pyconnectomist/pyconnectomist/scripts/pyconnectomist_preproc \
    -v 2 \
    -o /tmp/pyconnectomist \
    -s ab130187 \
    -i /neurospin/senior/nsap/data/V0/nifti/ab130187/000009_DTI/DTI.nii.gz \
    -r /neurospin/senior/nsap/data/V0/nifti/ab130187/000009_DTI/DTI.bvec \
    -b /neurospin/senior/nsap/data/V0/nifti/ab130187/000009_DTI/DTI.bval \
    -m Siemens \
    -t 2.46 \
    -f 0.75 \
    -a 2 \
    -u /neurospin/senior/nsap/data/V0/nifti/ab130187/000015_B0MAP/B0MAP.nii.gz \
    -p /neurospin/senior/nsap/data/V0/nifti/ab130187/000016_B0MAP/B0MAP.nii.gz \
    -x \
    -q 0.75 \
    -l 3.0 \
    -g /neurospin/senior/nsap/data/V0/morphologist \
    -e
"""


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


parser = argparse.ArgumentParser(description=doc,
                                 formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument(
    "-v", "--verbose", dest="verbose", type=int, choices=[0, 1, 2], default=0,
    help="increase the verbosity level: 0 silent, [1, 2] verbose.")
parser.add_argument(
    "-e", "--erase", dest="erase", action="store_true",
    help="if activated, clean the subject folder if already created.")
parser.add_argument(
    "-c", "--connectomistconfig", dest="connectomistconfig", metavar="PATH",
    help="the path to the Connectomist configuration file.", type=is_file)
parser.add_argument(
    "-o", "--outdir", dest="outdir", required=True, metavar="PATH",
    help="the Connectomist preprocessing home directory.", type=is_directory)
parser.add_argument(
    "-s", "--subjectid", dest="subjectid", required=True,
    help="the subject identifier.")
parser.add_argument(
    "-i", "--dwifile", dest="dwifile", metavar="FILE", required=True,
    help="the subject diffusion image to be processed.", type=is_file)
parser.add_argument(
    "-r", "--bvecfile", dest="bvecfile", metavar="FILE", required=True,
    help="the subject diffusion b-vectors to be processed.", type=is_file)
parser.add_argument(
    "-b", "--bvalfile", dest="bvalfile", metavar="FILE", required=True,
    help="the subject diffusion b-values to be processed.", type=is_file)
parser.add_argument(
    "-m", "--manufacturer", dest="manufacturer", choices=MANUFACTURERS.keys(),
    help="the MRI scanner manufacturer.")
parser.add_argument(
    "-t", "--delta_te", dest="delta_te", type=float, required=True,
    help=("the difference in msec between the 2 echoes of the B0 magnitude "
          "map."))
parser.add_argument(
    "-f", "--partial_fourier_factor", dest="partial_fourier_factor",
    type=float, required=True,
    help="the percentage of k-space plane acquired (]0;1]).")
parser.add_argument(
    "-a", "--parallel_acceleration_factor",
    dest="parallel_acceleration_factor", type=float, required=True,
    help="the number of parallel acquisition in the k-space plane.")
parser.add_argument(
    "-u", "--b0_magnitude", dest="b0_magnitude", metavar="FILE", required=True,
    help="the path to the B0 magnitude map (also contains phase for GE).",
    type=is_file)
parser.add_argument(
    "-p", "--b0_phase", dest="b0_phase", metavar="FILE",
    help="the path to the B0 phase map (not for GE).", type=is_file)
parser.add_argument(
    "-x", "--invertx", dest="invertx", action="store_true",
    help="if activated, invert the x-axis in diffusion model.")
parser.add_argument(
    "-y", "--inverty", dest="inverty", action="store_true",
    help="if activated, invert the y-axis in diffusion model.")
parser.add_argument(
    "-z", "--invertz", dest="invertz", action="store_true",
    help="if activated, invert the z-axis in diffusion model.")
parser.add_argument(
    "-n", "--negative_sign", dest="negative_sign",  action="store_true",
    help=("if activated, invert the direction of unwarping in the "
          "suceptibility distortion correction."))
parser.add_argument(
    "-q", "--echo_spacing", dest="echo_spacing", type=float, required=False,
    help=("the acquisition time in msec between 2 centers of 2 "
          "consecutively acquired lines in k-space (not for Philips)."))
parser.add_argument(
    "-k", "--epi_factor", dest="epi_factor", type=int, required=False,
    help="the number of echoes after one RF pulse, i.e. echo train length.")
parser.add_argument(
    "-l", "--b0_field", dest="b0_field", type=float, default=3.0,
    help="the B0 field intensity (Philips only).")
parser.add_argument(
    "-w", "--water_fat_shift", dest="water_fat_shift", type=float,
    default=4.68, help="Philips only, the spins frequency offsets.")
parser.add_argument(
    "-d", "--delete_steps", dest="delete_steps", action="store_true",
    help=("if activated, remove all intermediate files and directories at "
          "the end."))
parser.add_argument(
    "-g", "--morphologist_dir", dest="morphologist_dir", required=True,
    metavar="PATH", type=is_directory,
    help="the path to the morphologist processings home directory.")
args = parser.parse_args()


"""
First check if the Connectomist subject directory exists on the file system,
and clean it if requested.
"""
tool = "Connectomist preproc"
config = args.connectomistconfig or DEFAULT_CONNECTOMIST_PATH
tool_version = ConnectomistWrapper._connectomist_version_check(config)
if args.verbose > 0:
    print("[info] Start Connectomist preproc...")
    print("[info] Directory: {0}.".format(args.outdir))
    print("[info] Subject: {0}.".format(args.subjectid))
    print("[info] DWI: {0}.".format(args.dwifile))
    print("[info] BVEC: {0}.".format(args.bvecfile))
    print("[info] BVAL: {0}.".format(args.bvalfile))
    print("[info] Magnitude: {0}.".format(args.b0_magnitude))
    print("[info] Phase: {0}.".format(args.b0_phase))
outdir = args.outdir
subjectid = args.subjectid
subjdir = os.path.join(args.outdir, subjectid, "preproc")
dwi = args.dwifile
bvec = args.bvecfile
bval = args.bvalfile
manufacturer = args.manufacturer
delta_te = args.delta_te
partial_fourier_factor = args.partial_fourier_factor
parallel_acceleration_factor = args.parallel_acceleration_factor
b0_magnitude = args.b0_magnitude
b0_phase = args.b0_phase
invertx = args.invertx
inverty = args.inverty
invertz = args.invertz
negative_sign = args.negative_sign
echo_spacing = args.echo_spacing
epi_factor = args.epi_factor
b0_field = args.b0_field
water_fat_shift = args.water_fat_shift
delete_steps = args.delete_steps
morphologist_dir = args.morphologist_dir
inputs = [dwi, bvec, bval]
outputs = []
if not os.path.isdir(subjdir):
    os.makedirs(subjdir)
elif os.path.isdir(subjdir) and args.erase:
    shutil.rmtree(subjdir)
    os.mkdir(subjdir)


"""
Connectomist preproc: all steps
"""
returned_values = complete_preprocessing(
    subjdir,
    subjectid,
    dwi,
    bval,
    bvec,
    manufacturer,
    delta_te,
    partial_fourier_factor,
    parallel_acceleration_factor,
    b0_magnitude,
    b0_phase=b0_phase,
    invertX=invertx,
    invertY=inverty,
    invertZ=invertz,
    negative_sign=negative_sign,
    echo_spacing=echo_spacing,
    EPI_factor=epi_factor,
    b0_field=b0_field,
    water_fat_shift=water_fat_shift,
    delete_steps=delete_steps,
    morphologist_dir=morphologist_dir,
    path_connectomist=config)
preproc_dwi, preproc_bval, preproc_bvec, preproc_outliers = returned_values
outputs = [preproc_dwi, preproc_bval, preproc_bvec, preproc_outliers]
if args.verbose > 1:
    print("[result] In folder: {0}.".format(subjdir))
