#! /usr/bin/env python
##########################################################################
# NSAp - Copyright (C) CEA, 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 modules
from __future__ import print_function
import os
import shutil
import glob
import json
import argparse
from datetime import datetime
from pprint import pprint
import textwrap
import numpy
import nibabel


# Bredala module
try:
    import bredala
    bredala.USE_PROFILER = False
    bredala.register("pyconnectome.tractography.probabilist",
                     names=["probtrackx2"])
    bredala.register("pyconnectome.utils.filetools",
                     names=["surf2surf"])
    bredala.register("pyfreesurfer.conversions.volconvs",
                     names=["mri_binarize", "mri_convert"])
    bredala.register("pyconnectome.connectomes.complete",
                     names=["probtrackx2_connectome_complete"])
except:
    pass

# Package import
from pyconnectome import __version__ as version
from pyconnectome import DEFAULT_FSL_PATH
from pyconnectome.wrapper import FSLWrapper
from pyconnectome.connectomes.complete import probtrackx2_connectome_complete

# PyFreeSurfer import
from pyfreesurfer import DEFAULT_FREESURFER_PATH
from pyfreesurfer.wrapper import FSWrapper


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


# Command parameters
doc = """
FSL Probtrackx2
~~~~~~~~~~~~~~~

Perform FSL connectome using surface vertices as nodes of the network.

Perspectives
~~~~~~~~~~~~

Deal with different surfaces.

Command
~~~~~~~

python $HOME/git/pyfsl/pyconnectome/scripts/pyconnectome_probtrackx2_complete \
    -c /etc/fsl/5.0/fsl.sh \
    -s 102008 \
    -o /volatile/nsap/hcp \
    -b /neurospin/tmp/slefranc/HCP/102008/diffusion.bedpostX \
    -B /volatile/nsap/hcp/102008/nodif_brain.nii.gz \
    -M /volatile/nsap/hcp/102008/nodif_brain_mask.nii.gz \
    -S /neurospin/population/HCP/processed/freesurfer/102008/T1w \
    -a 1 \
    -i 11 \
    -L /volatile/nsap/hcp/102008/L.gii \
    -R /volatile/nsap/hcp/102008/R.gii \
    -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


parser = argparse.ArgumentParser(
    formatter_class=argparse.RawDescriptionHelpFormatter,
    description=textwrap.dedent(doc))
required = parser.add_argument_group("required arguments")
parser.add_argument(
    "-c", "--fslconfig", dest="fslconfig", metavar="<path>",
    help="the FSL configuration file.", type=is_file)
parser.add_argument(
    "-k", "--fsconfig", dest="fsconfig", metavar="<path>",
    help="the FreeSurfer configuration file.", type=is_file)
required.add_argument(
    "-s", "--subjectid", required=True, metavar="<id>",
    help="Subject identifier.")
required.add_argument(
    "-o", "--outdir", required=True, metavar="<path>",
    help="directory where to output.", type=is_directory)
required.add_argument(
    "-L", "--lhsurf", required=True, metavar="<path>",
    help="The left hemisphere surface.", type=is_file)
required.add_argument(
    "-R", "--rhsurf", required=True, metavar="<path>",
    help="The right hemisphere surface.", type=is_file)
required.add_argument(
    "-b", "--beddir", dest="bedpostxdir", required=True, metavar="<path>",
    help="the FSL bedpostx directory.", type=is_directory)
required.add_argument(
    "-B", "--nodifbrain", required=True, metavar="<path>",
    help="Preprocessed brain-only volume with bvalue=0.", type=is_file)
required.add_argument(
    "-M", "--nodifbrainmask", required=True, metavar="<path>",
    help="Brain binary mask of the diffusion sequence.", type=is_file,)
required.add_argument(
    "-S", "--subjectsdir", metavar="<path>", required=True,
    help="To FreeSurfer subjects home directory.", type=is_directory)
required.add_argument(
    "-i", "--index", dest="index", required=True, metavar="<id>",
    help="the index of the reference seed (in the first specified surface) to "
         "compute the corresponding profile.", type=int)
parser.add_argument(
    "-a", "--nsamples", dest="nsamples", metavar="<int>", default=5000,
    help="the number of samples in probtrackx.", type=int)
parser.add_argument(
    "-t", "--nsteps", dest="nsteps", metavar="<int>", default=2000,
    help="the number of steps per sample in probtrackx.", type=int)
parser.add_argument(
    "-l", "--steplength", dest="steplength", metavar="<float>", default=0.5,
    help="the propagation step in probtrackx.", type=float)
parser.add_argument(
    "-m", "--sampvox", dest="sampvox", metavar="<float>", default=2.0,
    help="random sampling sphere in probtrackx (in mm).", type=float)
parser.add_argument(
    "-e", "--erase", dest="erase", action="store_true",
    help="if activated, clean the subject folder.")
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.")
args = parser.parse_args()
inputs = vars(args)
verbose = inputs.pop("verbose")
fsl_config =  inputs.pop("fslconfig") or DEFAULT_FSL_PATH
freesurfer_config =  inputs.pop("fsconfig") or DEFAULT_FREESURFER_PATH


"""
First check if the subject FSL directory exists on the file system, and
clean it if requested. Check also the bedpostx directory.
"""
tool = "pyconnectome_probtrackx2_complete"
timestamp = datetime.now().isoformat()
tool_version = version
fsl_version = FSLWrapper([], shfile=fsl_config).version
freesurfer_version = FSWrapper([], freesurfer_config).version
params = locals()
runtime = dict([(name, params[name])
               for name in ("fsl_config", "tool", "tool_version",
                            "fsl_version", "timestamp", "freesurfer_config",
                            "freesurfer_version")])
outputs = None
if verbose > 0:
    print("[info] Starting FSL probtrackx2 ...")
    print("[info] Runtime:")
    pprint(runtime)
    print("[info] Inputs:")
    pprint(inputs)
subjectdir = os.path.join(inputs["outdir"], inputs["subjectid"])
wdir = os.path.join(subjectdir, "{0}".format(inputs["index"]))
if inputs["erase"] and os.path.isdir(subjectdir):
    shutil.rmtree(subjectdir)
if not os.path.isdir(subjectdir):
    os.mkdir(subjectdir)
merged_prefix = os.path.join(inputs["bedpostxdir"], "merged")
merged_files = glob.glob(merged_prefix + "*")
if len(merged_files) == 0:
    raise ValueError("'{0}' is not a valid FSL bedpostx folder.".format(
        inputs["bedpostxdir"]))
if not os.path.isdir(wdir):
    os.mkdir(wdir)


"""
Start the connectome computation.
"""
coords, weights = probtrackx2_connectome_complete(
    outdir=wdir,
    subject_id=inputs["subjectid"],
    lh_surf=inputs["lhsurf"],
    rh_surf=inputs["rhsurf"],
    nodif_brain=inputs["nodifbrain"],
    nodif_brain_mask=inputs["nodifbrainmask"],
    bedpostx_dir=inputs["bedpostxdir"],
    nsamples=inputs["nsamples"],
    nsteps=inputs["nsteps"],
    steplength=inputs["steplength"],
    subjects_dir=inputs["subjectsdir"],
    loopcheck=True,
    cthr=0.2,
    fibthresh=0.01,
    distthresh=0.0,
    sampvox=inputs["sampvox"],
    fs_sh=freesurfer_config,
    fsl_sh=fsl_config)

"""
Update the outputs and save them and the inputs in a 'logs' directory.
"""
logdir = os.path.join(wdir, "logs")
if not os.path.isdir(logdir):
    os.mkdir(logdir)
shutil.move(os.path.join(wdir, "probtrackx.log"),
            os.path.join(logdir, "probtrackx.log"))
params = locals()
outputs = dict([(name, params[name])
                for name in ("coords", "weights")])
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)

