#! /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 os
import shutil
import argparse
from datetime import datetime
import json
from pprint import pprint
import nibabel
import numpy
import textwrap
from argparse import RawTextHelpFormatter

# Bredala import
try:
    import bredala
    bredala.USE_PROFILER = False
    bredala.register("pyfreesurfer.conversions.volconvs",
                     names=["mri_convert"])
    bredala.register("pyfreesurfer.conversions.surfconvs",
                     names=["resample_cortical_surface", "surf_convert",
                            "midgray_surface", "interhemi_surfreg",
                            "measure_smoothing"])
    bredala.register("pyfreesurfer.utils.regtools",
                     names=["conformed_to_native_space"])
except:
    pass

# Pyfreesurfer import
from pyfreesurfer import __version__ as version
from pyfreesurfer.wrapper import FSWrapper
from pyfreesurfer import DEFAULT_FREESURFER_PATH
from pyfreesurfer import DEFAULT_TEMPLATE_SYM_PATH
from pyfreesurfer.conversions.volconvs import mri_convert
from pyfreesurfer.conversions.surfconvs import surf_convert
from pyfreesurfer.conversions.surfconvs import resample_cortical_surface
from pyfreesurfer.conversions.surfconvs import midgray_surface
from pyfreesurfer.conversions.surfconvs import interhemi_surfreg
from pyfreesurfer.conversions.surfconvs import measure_smoothing
from pyfreesurfer.utils.regtools import conformed_to_native_space


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


# Script documentation
DOC = """
Freesurfer conversion
~~~~~~~~~~~~~~~~~~~~~

Convert the results returned by the FreeSurfer cortical reconstruction
pipeline.

Steps:

1- Nifti conversions: aseg - aparc+aseg - aparc.a2009s+aseg - wm - t1.
   Export FreeSurfer '.mgz' images of interest in Nifti format. These
   images are resliced like the 'rawavg.mgz' file, have a '.native'
   suffix and are stored in a 'convert' folder.

2- Registration matrix: between the conformed space (orig.mgz)
   and the native anatomical (rawavg.mgz).

3- Create the ribbon right and left hemispheres masks.

4- Create the mid thickness white and pial, right and left surfaces.
   (optional 10-30 min)

5- Surface conversions: resample the white or pial FreeSurfer
   surfaces at different resolutions (impacts the number of vertex)
   with common mesh that can be directly used in a longitudinal
   setting. The results are also stored in a 'convert' folder with
   a '.native' suffix and the considered level in the file name. Vetex
   are expressed in the index coordinate system.

6- Align the interhemispheric surface vertices by aplying an existing atlas.
   [optional 1-2 hours]

7- Smooth the thickness texture with a fwhm of 10 (optional)

Command:

python $HOME/git/pyfreesurfer/pyfreesurfer/scripts/pyfreesurfer_conversion \
    -v 2 \
    -c /i2bm/local/freesurfer/SetUpFreeSurfer.sh \
    -d /neurospin/senior/nsap/data/V4/freesurfer \
    -s ag110371 \
    -o /tmp \
    -r lhrh \
    -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


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 pyfreesurfer_conversion",
        description=textwrap.dedent(DOC),
        formatter_class=RawTextHelpFormatter)

    # Required arguments
    required = parser.add_argument_group("required arguments")
    required.add_argument(
        "-d", "--fsdir",
        required=True, metavar="PATH", type=is_directory,
        help="the FreeSurfer processing home directory.")
    required.add_argument(
        "-o", "--outdir",
        required=True, metavar="PATH", type=is_directory,
        help="the FreeSurfer conversion home directory.")
    parser.add_argument(
        "-s", "--subjectid",
        required=True,
        help="the subject identifier.")

    # Optional arguments
    parser.add_argument(
        "-v", "--verbose",
        type=int, choices=[0, 1, 2], default=0,
        help="increase the verbosity level: 0 silent, [1, 2] verbose.")
    parser.add_argument(
        "-e", "--erase",
        action="store_true",
        help="if activated, clean the result folder.")
    parser.add_argument(
        "-c", "--config", dest="fsconfig",
        metavar="FILE", type=is_file,
        help="the FreeSurfer configuration file.")
    parser.add_argument(
        "-m", "--midgray",
        action="store_true",
        help=("if activated, compute the mid gray thickness surfaces."))
    parser.add_argument(
        "-f", "--smooth",
        action="store_true",
        help=("if activated, smooth the thickness texture with a fwhm of 10."))
    parser.add_argument(
        "-r", "--surfreg",
        choices=["lh", "rh", "lhrh"],
        help=("if set, surface-based interhemispheric registration using the "
              "specified hemisphere."))
    parser.add_argument(
        "-t", "--templatesym",
        type=is_directory,
        help=("path to the 'fsaverage_sym' template."))

    # Create a dict of arguments to pass to the 'main' function
    args = parser.parse_args()
    if args.fsconfig is None:
        args.fsconfig = DEFAULT_FREESURFER_PATH

    return args


"""
Parse the command line.
"""
args = get_cmd_line_args()
tool = "pyfreesurfer_conversion"
timestamp = datetime.now().isoformat()
tool_version = version
freesurfer_config = args.fsconfig
freesurfer_version = FSWrapper([], freesurfer_config).version
params = locals()
runtime = dict([(name, params[name])
               for name in ("freesurfer_config", "tool", "tool_version",
                            "freesurfer_version", "timestamp")])
if args.verbose > 0:
    print("[info] Start FreeSurfer conversions...")
    print("[info] FreeSurfer home directory: {0}.".format(args.fsdir))
    print("[info] Output directory: {0}.".format(args.outdir))
    print("[info] Subject: {0}.".format(args.subjectid))
    print("[info] FreeSurfer version: {0}.".format(freesurfer_version))
subjectid = args.subjectid
templatesym = args.templatesym or DEFAULT_TEMPLATE_SYM_PATH
subjdir = os.path.join(args.fsdir, subjectid)
convertdir = os.path.join(args.outdir, args.subjectid, "convert")
imagedir = os.path.join(convertdir, "images")
surfacedir = os.path.join(convertdir, "surfaces")
texturedir = os.path.join(convertdir, "textures")
addmidgray = args.midgray
surfreg = args.surfreg
smooth = args.smooth
params = locals()
inputs = dict([(name, params[name])
               for name in ("subjectid", "subjdir", "convertdir", "imagedir",
                            "surfacedir", "addmidgray", "surfreg",
                            "templatesym", "smooth", "texturedir")])
outputs = None
if not os.path.isdir(subjdir):
    raise ValueError(
        "'{0}' is not a FreeSurfer subject folder.".format(subjdir))
if args.erase and os.path.isdir(convertdir):
    shutil.rmtree(convertdir)
for directory in (convertdir, imagedir, surfacedir, texturedir):
    if not os.path.isdir(directory):
        os.makedirs(directory)


"""
Step 1: Nifti conversions.
"""
if args.verbose > 0:
    print("[info] Start Nifti conversions...")
niftifiles = {}
for modality in ["aparc+aseg", "aparc.a2009s+aseg", "aseg", "wm", "rawavg",
                 "ribbon", "brain"]:
    regex = os.path.join(args.subjectid, "mri", "{0}.mgz".format(modality))
    niftifiles[modality] = mri_convert(
        args.fsdir,
        regex,
        args.outdir,
        destdirname=os.path.join("convert", "images"),
        reslice=True,
        interpolation="nearest",
        fsconfig=freesurfer_config)
    if args.verbose > 1:
        print("[result] {0}: {1}.".format(modality, niftifiles[modality]))


"""
Step 2: Registration matrix.
"""
if args.verbose > 0:
    print("[info] Start Registration matrix...")
regex = os.path.join(args.subjectid, "mri")
trffile = conformed_to_native_space(
    args.fsdir,
    regex,
    args.outdir,
    fsconfig=freesurfer_config)
if args.verbose > 1:
    print("[result] trffile: {0}.".format(trffile))


"""
Step 3: Create the ribbon right and left hemispheres masks.
"""
ribbon_file = niftifiles["ribbon"][0]
image = nibabel.load(ribbon_file)
data = image.get_data()
for hemi, label in [("lh", 3), ("rh", 42)]:
    indices = numpy.where((data > (label - 0.01)) & (data < (label + 0.01)))
    mask = numpy.zeros(data.shape, dtype=int)
    mask[indices] = 1
    mask_image = nibabel.Nifti1Image(mask, image.get_affine())
    mask_file = os.path.join(imagedir, "{0}.ribbon.nii.gz".format(hemi))
    nibabel.save(mask_image, mask_file)
    niftifiles["ribbon"].append(mask_file)


"""
Step 4: Create the mid thickness right and left hemispheres surfaces
"""
if addmidgray:
    for hemi in ["lh", "rh"]:
        surf = midgray_surface(
            hemi,
            surfacedir,
            args.fsdir,
            subjectid,
            fsconfig=freesurfer_config)


"""
Step 5: Surface conversions.
"""
if args.verbose > 0:
    print("[info] Start surface conversions...")
surfaces = {}
annotations = []
modalities = ["pial", "white"]
if addmidgray:
    modalities.append("graymid")
for modality in modalities:
    for hemi in ["lh", "rh"]:
        name = "{0}.{1}".format(hemi, modality)
        regex = os.path.join(args.subjectid, "surf", name)
        resamplefiles, annotfiles = resample_cortical_surface(
            args.fsdir,
            regex,
            args.outdir,
            destdirname=os.path.join("convert", "surfaces"),
            orders=[4, 5, 6, 7],
            surface_name=modality,
            fsconfig=freesurfer_config)
        annotations.extend(annotfiles)
        surfaces[name] = surf_convert(
            args.fsdir,
            niftifiles["rawavg"],
            resamplefiles,
            sidpos=-4,
            rm_orig=True,
            fsconfig=freesurfer_config)
        if args.verbose > 1:
            print("[result] {0}: {1}.".format(name, surfaces[name]))
annotations = list(set(annotations))
if args.verbose > 1:
    print("[result] Annotations: {0}.".format(annotations))


"""
Step 6: Surface-based interhemispheric registration.
"""
xhemidirs = {}
spherefiles = {}
if surfreg is not None:
    if surfreg == "lhrh":
        hemis = ["lh", "rh"]
    else:
        hemis = [surfreg]
    for hemi in hemis:
        xhemidir, spherefile = interhemi_surfreg(
                hemi=hemi,
                outdir=convertdir,
                fsdir=args.fsdir,
                sid=subjectid,
                template_file=templatesym,
                destname="{0}_surfreg".format(hemi),
                fsconfig=freesurfer_config)
        xhemidirs[hemi] = xhemidir
        spherefiles[hemi] = spherefile
        if args.verbose > 1:
            print("[result] xhemidir: {0}.".format(xhemidir))
            print("[result] spherefile: {0}.".format(spherefile))


"""
Step 7: Smooth the thickness texture with a fwhm of 10.
Note that we need an actual real subject and use fsaverage that is the same as
ico7.
"""
if smooth:
    resampled_files, smoothed_files = measure_smoothing(
        measure="thickness",
        fsdir=args.fsdir,
        sid=subjectid,
        outdir=texturedir,
        fwhm=10,
        fsconfig=freesurfer_config)
else:
    resampled_files, smoothed_files = (None, None)


"""
Update the outputs and save them and the inputs in a 'logs' directory.
"""
logdir = os.path.join(convertdir, "logs")
if not os.path.isdir(logdir):
    os.mkdir(logdir)
params = locals()
outputs = dict([(name, params[name])
               for name in ("niftifiles", "trffile", "surfaces",
                            "annotations", "xhemidirs", "spherefiles",
                            "resampled_files", "smoothed_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 args.verbose > 1:
    print("[final]")
    pprint(outputs)

