#! /usr/bin/env python
# -*- coding: utf-8 -*
##########################################################################
# 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 import
import os
import argparse
import json
from pprint import pprint
from datetime import datetime
import textwrap
from argparse import RawTextHelpFormatter

# Bredala module
try:
    import bredala
    bredala.USE_PROFILER = False
    bredala.register("pyconnectome.connectomes.reduced",
                     names=["probtrackx2_connectome"])
except:
    pass

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

# Package import
from pyconnectome import __version__ as version
from pyconnectome import DEFAULT_FSL_PATH
from pyconnectome.wrapper import FSLWrapper
from pyconnectome.connectomes.reduced import probtrackx2_connectome


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


DOC = """
Compute the connectome of a given parcellation, like the FreeSurfer aparc+aseg
segmentation, using FSL Probtrackx2. Assuming you have run Freesurfer and
bedpostx.

Example of command for HCP, propagation through the white-grey interface:

python ~/git/pyconnectome/pyconnectome/scripts/pyconnectome_reduced_probtrackx2 \
    -s ${subject} \
    -B ${root_dir}/dwi_preproc_wo_fieldmap/${subject}/nodif_brain.nii.gz \
    -m ${root_dir}/dwi_preproc_wo_fieldmap/${subject}/nodif_brain_mask.nii.gz \
    -x ${root_dir}/bedpostx_wo_fieldmap/${subject}/diffusion.bedpostX \
    -p ${root_dir}/freesurfer/${subject}/mri/aparc+aseg.mgz \
    -f ${FREESURFER_HOME}/FreeSurferColorLUT.txt \
    -c lausanne2008 \
    -N 500 \
    -P 200 \
    -L 1.25 \
    -o /volatile/imagen/probtrackx2/${subject} \
    -d /volatile/imagen/probtrackx2/tmp \
    -u \
    -S /neurospin/imagen/BL/processed/freesurfer \
    -v 2

Example of command for HCP, propagation through the white matter with fiber
length normalization:
"""


def is_file(filepath):
    """ Check file's existence - argparse 'type' argument.
    """
    if not os.path.isfile(filepath):
        raise argparse.ArgumentError("File does not exist: %s" % filepath)
    return filepath


def is_dir(dirpath):
    """ Check direcory's existence - argparse 'type' argument.
    """
    if not os.path.isdir(dirpath):
        raise argparse.ArgumentError("Directory does not exist: %s" % dirpath)
    return dirpath


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

    # Required arguments
    required = parser.add_argument_group("required arguments")
    required.add_argument(
        "-s", "--subject-id",
        required=True, metavar="<id>",
        help="Subject identifier.")
    required.add_argument(
        "-b", "--nodif-brain",
        type=is_file, required=True, metavar="<path>",
        help="Preprocessed brain-only volume with bvalue=0.")
    required.add_argument(
        "-m", "--nodif-brain-mask",
        type=is_file, required=True, metavar="<path>",
        help="Brain binary mask in diffusion.")
    required.add_argument(
        "-x", "--bedpostx-dir",
        type=is_dir, required=True, metavar="<path>",
        help="Bedpostx output directory.")
    required.add_argument(
        "-p", "--t1-parc",
        type=is_file, required=True, metavar="<path>",
        help="Path to the parcellation that defines the nodes "
             "of the connectome, e.g. aparc+aseg.mgz from Free"
             "surfer. Has to be in Freesurfer T1 space.")
    required.add_argument(
        "-f", "--t1-parc-lut",
        type=is_file, required=True, metavar="<path>",
        help="Path to the Look Up Table for the passed parcellation in the "
             "Freesurfer LUT format. If your T1 parcellation comes from "
             "Freesurfer, this will most likely be "
             "<$FREESURFER_HOME>/FreeSurferColorLUT.txt.")
    required.add_argument(
        "-c", "--connectome-lut",
        metavar="<path>", required=True,
        help="Path to a Look Up Table in the FreeSurfer LUT format, listing "
             "the regions from the parcellation to use as nodes in the connec"
             "tome. The region names should match the ones used in the "
             "<t1_parc_lut> LUT and the integer labels should be the row/col "
             "positions in the connectome. Alternatively it can be set to "
             "'Lausanne2008' to use the predefined LUT for the Lausanne 2008 a"
             "tlas, which is based on the FreeSurfer aparc+aseg parcellation.")
    required.add_argument(
        "-n", "--nsamples",
        type=int, required=True, metavar="<int>",
        help="Number of samples per voxel to initiate in seed region")
    required.add_argument(
        "-e", "--nsteps",
        type=int, required=True, metavar="<int>",
        help="Maximum number of steps for a sample.")
    required.add_argument(
        "-l", "--steplength",
        type=float, required=True, metavar="<float>",
        help="Step length in mm.")
    required.add_argument(
        "-o", "--outdir",
        required=True, metavar="<path>",
        help="Directory where to output.")
    required.add_argument(
        "-d", "--tempdir",
        required=True, metavar="<path>",
        help="Where to write temporary directories e.g. /tmp.")

    # Optional arguments
    parser.add_argument(
        "-U", "--fix-freesurfer-subcortical",
        action="store_true",
        help="If the <t1_parc> is aparc+aseg or aparc.a2009s+aseg from "
             "FreeSurfer set this option to True, to recompute the "
             "subcortical segmentations of 5 structures that are uncorrectly "
             "segmented by FreeSurfer, using FSL FIRST")
    parser.add_argument(
        "-S", "--subjects-dir",
        metavar="<path>",
        help="To set or bypass the $SUBJECTS_DIR environment variable.")
    parser.add_argument(
        "-Q", "--no-snapshots",
        action="store_false", dest="snapshots", default=True,
        help="To not create snapshots of the results. "
             "Useful if matplotlib is not installed.")
    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(
        "-v", "--verbose",
        type=int, choices=[0, 1, 2], default=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 = "pyconnectome_reduced_probtrackx2"
timestamp = datetime.now().isoformat()
tool_version = version
fsl_version = FSLWrapper([], shfile=inputs["fsl_sh"]).version
freesurfer_version = FSWrapper([], inputs["fs_sh"]).version
params = locals()
runtime = dict([(name, params[name])
               for name in ("tool", "tool_version", "fsl_version",
                            "timestamp", "freesurfer_version")])
outputs = None
if verbose > 0:
    pprint(runtime)
    pprint(inputs)


"""
Start the connectome computation.
"""
connectome_file, labels_file, connectome_snap_file = probtrackx2_connectome(
    **inputs)


"""
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 ("connectome_file", "labels_file",
                            "connectome_snap_file")])
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:
    pprint("[info] Outputs:")
    pprint(outputs)
