#! /usr/bin/env python
# -*- coding: utf-8 -*
##########################################################################
# NSAp - Copyright (C) CEA, 2017
# 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.
##########################################################################

# Standard
import os
import argparse
import datetime
import json
from pprint import pprint
import textwrap
from argparse import RawTextHelpFormatter

# Package
import pyfreesurfer
from pyfreesurfer import DEFAULT_FREESURFER_PATH
from pyfreesurfer.segmentation.cortical import recon_all_longitudinal
from pyfreesurfer.wrapper import FSWrapper


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


# Script documentation
DOC = """
Assuming you have run recon-all for all timepoints of a given subject,
and that the results are stored in one SUBJECTS_DIR per timepoint, this
script will:
    - create a template for the subject and process it with recon-all
    - rerun recon-all for all timepoints of the subject using the template
"""


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 get_cmd_line_args():
    """
    Create a command line argument parser and return a dict mapping
    <argument name> -> <argument value>.
    """
    usage = ("%(prog)s -o <dir> -s <subject id> -S <dir> <dir> ... "
             "[-t <name> <name> ...] [-c <path>] [-v <int>]")
    prog = "python pyfreesurfer_reconall_longitudinal"
    parser = argparse.ArgumentParser(
        prog=prog,
        usage=usage,
        description=textwrap.dedent(DOC),
        formatter_class=RawTextHelpFormatter)

    # Required arguments
    parser.add_argument("-o", "--outdir", required=True, metavar="<path>",
                        help="Directory where to output. Created if not "
                             "already existing.")
    parser.add_argument("-s", "--subject-id", required=True, metavar="<id>",
                        help="Identifier of subject, used for all timepoints.")
    parser.add_argument("-S", "--subjects-dirs", required=True, nargs="+",
                        metavar="<path> <path> ...",
                        help="The FreeSurfer SUBJECTS_DIRs of timepoints.")

    # Optional arguments
    thelp = ("The timepoint names in the same order as the SUBJECTS_DIRs. "
             "Used to create the subject longitudinal IDs. "
             'By default timepoints are "1", "2"...')
    parser.add_argument("-t", "--timepoints", nargs="+",
                        metavar="<name> <name> ...", help=thelp)
    parser.add_argument("-c", "--config", type=is_file, metavar="<path>",
                        dest="fsconfig",
                        help="the FreeSurfer configuration file. "
                             "By default %s." % DEFAULT_FREESURFER_PATH)
    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["fsconfig"] is None:
        kwargs["fsconfig"] = DEFAULT_FREESURFER_PATH

    return kwargs, verbose


# Parse the command line.
inputs, verbose = get_cmd_line_args()

# Runtime informations
runtime = dict(tool="pyfreesurfer_reconall_longitudinal",
               tool_version=pyfreesurfer.__version__,
               fs_version=FSWrapper([], shfile=inputs["fsconfig"]).version,
               timestamp=datetime.datetime.now().isoformat())

if verbose > 0:
    pprint("[info] Starting longitudinal recon-all")
    pprint("[info] Runtime:")
    pprint(runtime)
    pprint("[info] Inputs:")
    pprint(inputs)

# Run
subject_template_id, subject_long_ids = recon_all_longitudinal(**inputs)

###############################################################################
# Create a 'logs' directory in each recon-all result to store inputs and
# outputs as JSONs

# Subject template recon-all
subject_template_dir = os.path.join(inputs["outdir"], subject_template_id)
template_logs_dir = os.path.join(subject_template_dir, "logs")
if not os.path.isdir(template_logs_dir):
    os.mkdir(template_logs_dir)
# Write JSONs
for k, v in dict(inputs=inputs, outputs=dict(), runtime=runtime).items():
    path_json = os.path.join(template_logs_dir, "%s.json" % k)
    with open(path_json, "w") as f:
        json.dump(v, f, sort_keys=True, check_circular=True, indent=4)

# Subject longitudinal recon-alls
for subject_long_id, timepoint, subjects_dir in zip(subject_long_ids,
                                                    inputs["timepoints"],
                                                    inputs["subjects_dirs"]):
    subject_long_dir = os.path.join(inputs["outdir"], subject_long_id)
    long_logs_dir = os.path.join(subject_long_dir, "logs")
    if not os.path.isdir(long_logs_dir):
        os.mkdir(long_logs_dir)
    long_inputs = dict(subject_id=inputs["subject_id"],
                       subjects_dir=subjects_dir,
                       timepoint=timepoint,
                       outdir=inputs["outdir"],
                       subject_template_id=subject_template_id,
                       subject_template_dir=subject_template_dir)
    for k, v in dict(inputs=long_inputs, outputs=dict(), runtime=runtime).items():
        path_json = os.path.join(template_logs_dir, "%s.json" % k)
        with open(path_json, "w") as f:
            json.dump(v, f, sort_keys=True, check_circular=True, indent=4)

# Global
outputs = dict(subject_template_id=subject_template_id,
               subject_long_ids=subject_long_ids)
