#! /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.tractography",
                     names=["complete_tractography"])
    bredala.register("pyconnectomist.tractography.model",
                     names=["dwi_local_modeling", "export_scalars_to_nifti"])
    bredala.register("pyconnectomist.wrappers",
                     names=["ConnectomistWrapper.__call__"])
except:
    pass

# Clindmri import
from pyconnectomist import __version__ as version
from pyconnectomist.tractography import complete_tractography
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", "preprocdir", "morphologistdir"]

# Script documentation
doc = """
Connectomist DTI fit
~~~~~~~~~~~~~~~~~~~~

Function that runs all the Connectomist 'Local modeling' tab to estimate a DTI
model from the preprocessed diffusion data.
Generates results in '<outdir>/<subjectid>/dtifit'.

Steps:

1- Create the model output directory if not existing.
2- Detect the Connectomist eddy motion correction folder.
3- Detect the Connectomist rough mask folder.
4- Compute the second order tensor diffusion model.
5 - Export diffusion scalars.

Command:

python $HOME/git/pyconnectomist/pyconnectomist/scripts/pyconnectomist_dtifit \
    -v 2 \
    -o /tmp/pyconnectomist \
    -g /neurospin/senior/nsap/data/V0/morphologist \
    -s ab130187 \
    -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 tractography home directory.", type=is_directory)
parser.add_argument(
    "-s", "--subjectid", dest="subjectid", required=True,
    help="the subject identifier.")
parser.add_argument(
    "-p", "--preprocdir", dest="preprocdir", metavar="PATH",
    help=("the path to the Connectomist diffusion preprocessings. If not "
          "specified assume data in '<outdir>/<subjectid>/preproc'."),
    type=is_directory)
parser.add_argument(
    "-g", "--morphologistdir", dest="morphologistdir", 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 tract"
config = args.connectomistconfig or DEFAULT_CONNECTOMIST_PATH
tool_version = ConnectomistWrapper._connectomist_version_check(config)
if args.verbose > 0:
    print("[info] Start Connectomist DTI fit...")
    print("[info] Directory: {0}.".format(args.outdir))
    print("[info] Subject: {0}.".format(args.subjectid))
    print("[info] Preproc: {0}.".format(args.preprocdir))
outdir = args.outdir
subjectid = args.subjectid
subjdir = os.path.join(args.outdir, subjectid, "dtifit")
preprocdir = args.preprocdir
morphologistdir = args.morphologistdir
if preprocdir is None:
    preprocdir = os.path.join(args.outdir, subjectid, "preproc")
    if args.verbose > 0:
        print("[info] Preproc: {0}.".format(preprocdir))
if not os.path.isdir(subjdir):
    os.makedirs(subjdir)
elif os.path.isdir(subjdir) and args.erase:
    shutil.rmtree(subjdir)
    os.mkdir(subjdir)
inputs = [preprocdir]
outputs = []

"""
Connectomist tractography: DTI fit only
"""
gfa, md, _, _ = complete_tractography(
    subjdir,
    preprocdir,
    morphologistdir,
    subjectid,
    model="dti",
    rgbscale=1.0,
    model_only=True,
    path_connectomist=config)
outputs = [gfa, md]
if args.verbose > 1:
    print("[result] In folder: {0}.".format(subjdir))
