# -*- coding: utf-8 -*-
# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
"""Provides interfaces to various longitudinal commands provided by freesurfer
"""
from __future__ import (print_function, division, unicode_literals,
                        absolute_import)

import os
import os.path

from ... import logging
from ...utils.filemanip import split_filename, copyfile

from .base import (FSCommand, FSTraitedSpec, FSScriptCommand,
                   FSScriptOutputSpec, FSCommandOpenMP, FSTraitedSpecOpenMP)
from ..base import (isdefined, TraitedSpec, File, traits, Directory)

__docformat__ = 'restructuredtext'
iflogger = logging.getLogger('nipype.interface')


class MPRtoMNI305InputSpec(FSTraitedSpec):
    # environment variables, required
    # usedefault=True is hack for on_trait_change in __init__
    reference_dir = Directory(
        "", exists=True, mandatory=True, usedefault=True, desc="TODO")
    target = traits.String(
        "", mandatory=True, usedefault=True, desc="input atlas file")
    # required
    in_file = File(
        argstr='%s',
        usedefault=True,
        desc="the input file prefix for MPRtoMNI305")


class MPRtoMNI305OutputSpec(FSScriptOutputSpec):
    out_file = File(
        exists=False,
        desc="The output file '<in_file>_to_<target>_t4_vox2vox.txt'")


class MPRtoMNI305(FSScriptCommand):
    """
    For complete details, see FreeSurfer documentation

    Examples
    ========

    >>> from nipype.interfaces.freesurfer import MPRtoMNI305, Info
    >>> mprtomni305 = MPRtoMNI305()
    >>> mprtomni305.inputs.target = 'structural.nii'
    >>> mprtomni305.inputs.reference_dir = '.' # doctest: +SKIP
    >>> mprtomni305.cmdline # doctest: +SKIP
    'mpr2mni305 output'
    >>> mprtomni305.inputs.out_file = 'struct_out' # doctest: +SKIP
    >>> mprtomni305.cmdline # doctest: +SKIP
    'mpr2mni305 struct_out' # doctest: +SKIP
    >>> mprtomni305.inputs.environ['REFDIR'] == os.path.join(Info.home(), 'average') # doctest: +SKIP
    True
    >>> mprtomni305.inputs.environ['MPR2MNI305_TARGET'] # doctest: +SKIP
    'structural'
    >>> mprtomni305.run() # doctest: +SKIP

    """
    _cmd = "mpr2mni305"
    input_spec = MPRtoMNI305InputSpec
    output_spec = MPRtoMNI305OutputSpec

    def __init__(self, **inputs):
        super(MPRtoMNI305, self).__init__(**inputs)
        self.inputs.on_trait_change(self._environ_update, 'target')
        self.inputs.on_trait_change(self._environ_update, 'reference_dir')

    def _format_arg(self, opt, spec, val):
        if opt in ['target', 'reference_dir']:
            return ""
        elif opt == 'in_file':
            _, retval, ext = split_filename(val)
            # Need to copy file to working cache directory!
            copyfile(
                val,
                os.path.abspath(retval + ext),
                copy=True,
                hashmethod='content')
            return retval
        return super(MPRtoMNI305, self)._format_arg(opt, spec, val)

    def _environ_update(self):
        # refdir = os.path.join(Info.home(), val)
        refdir = self.inputs.reference_dir
        target = self.inputs.target
        self.inputs.environ['MPR2MNI305_TARGET'] = target
        self.inputs.environ["REFDIR"] = refdir

    def _get_fname(self, fname):
        return split_filename(fname)[1]

    def _list_outputs(self):
        outputs = super(MPRtoMNI305, self)._list_outputs()
        fullname = "_".join([
            self._get_fname(self.inputs.in_file), "to", self.inputs.target,
            "t4", "vox2vox.txt"
        ])
        outputs['out_file'] = os.path.abspath(fullname)
        return outputs


class RegisterAVItoTalairachInputSpec(FSTraitedSpec):
    in_file = File(
        argstr='%s',
        exists=True,
        mandatory=True,
        position=0,
        desc="The input file")
    target = File(
        argstr='%s',
        exists=True,
        mandatory=True,
        position=1,
        desc="The target file")
    vox2vox = File(
        argstr='%s',
        exists=True,
        mandatory=True,
        position=2,
        desc="The vox2vox file")
    out_file = File(
        'talairach.auto.xfm',
        usedefault=True,
        argstr='%s',
        position=3,
        desc="The transform output")


class RegisterAVItoTalairachOutputSpec(FSScriptOutputSpec):
    out_file = File(
        exists=False, desc="The output file for RegisterAVItoTalairach")


class RegisterAVItoTalairach(FSScriptCommand):
    """
    converts the vox2vox from talairach_avi to a talairach.xfm file

    This is a script that converts the vox2vox from talairach_avi to a
    talairach.xfm file. It is meant to replace the following cmd line:

    tkregister2_cmdl \
        --mov $InVol \
        --targ $FREESURFER_HOME/average/mni305.cor.mgz \
        --xfmout ${XFM} \
        --vox2vox talsrcimg_to_${target}_t4_vox2vox.txt \
        --noedit \
        --reg talsrcimg.reg.tmp.dat
    set targ = $FREESURFER_HOME/average/mni305.cor.mgz
    set subject = mgh-02407836-v2
    set InVol = $SUBJECTS_DIR/$subject/mri/orig.mgz
    set vox2vox = $SUBJECTS_DIR/$subject/mri/transforms/talsrcimg_to_711-2C_as_mni_average_305_t4_vox2vox.txt

    Examples
    ========

    >>> from nipype.interfaces.freesurfer import RegisterAVItoTalairach
    >>> register = RegisterAVItoTalairach()
    >>> register.inputs.in_file = 'structural.mgz'                         # doctest: +SKIP
    >>> register.inputs.target = 'mni305.cor.mgz'                          # doctest: +SKIP
    >>> register.inputs.vox2vox = 'talsrcimg_to_structural_t4_vox2vox.txt' # doctest: +SKIP
    >>> register.cmdline                                                   # doctest: +SKIP
    'avi2talxfm structural.mgz mni305.cor.mgz talsrcimg_to_structural_t4_vox2vox.txt talairach.auto.xfm'

    >>> register.run() # doctest: +SKIP
    """
    _cmd = "avi2talxfm"
    input_spec = RegisterAVItoTalairachInputSpec
    output_spec = RegisterAVItoTalairachOutputSpec

    def _list_outputs(self):
        outputs = self.output_spec().get()
        outputs['out_file'] = os.path.abspath(self.inputs.out_file)
        return outputs


class EMRegisterInputSpec(FSTraitedSpecOpenMP):
    # required
    in_file = File(
        argstr="%s",
        exists=True,
        mandatory=True,
        position=-3,
        desc="in brain volume")
    template = File(
        argstr="%s",
        exists=True,
        mandatory=True,
        position=-2,
        desc="template gca")
    out_file = File(
        argstr="%s",
        exists=False,
        name_source=['in_file'],
        name_template="%s_transform.lta",
        hash_files=False,
        keep_extension=False,
        position=-1,
        desc="output transform")
    # optional
    skull = traits.Bool(
        argstr="-skull", desc="align to atlas containing skull (uns=5)")
    mask = File(argstr="-mask %s", exists=True, desc="use volume as a mask")
    nbrspacing = traits.Int(
        argstr="-uns %d",
        desc=
        "align to atlas containing skull setting unknown_nbr_spacing = nbrspacing"
    )
    transform = File(
        argstr="-t %s", exists=True, desc="Previously computed transform")


class EMRegisterOutputSpec(TraitedSpec):
    out_file = File(exists=False, desc="output transform")


class EMRegister(FSCommandOpenMP):
    """ This program creates a tranform in lta format

    Examples
    ========
    >>> from nipype.interfaces.freesurfer import EMRegister
    >>> register = EMRegister()
    >>> register.inputs.in_file = 'norm.mgz'
    >>> register.inputs.template = 'aseg.mgz'
    >>> register.inputs.out_file = 'norm_transform.lta'
    >>> register.inputs.skull = True
    >>> register.inputs.nbrspacing = 9
    >>> register.cmdline
    'mri_em_register -uns 9 -skull norm.mgz aseg.mgz norm_transform.lta'
    """
    _cmd = 'mri_em_register'
    input_spec = EMRegisterInputSpec
    output_spec = EMRegisterOutputSpec

    def _list_outputs(self):
        outputs = self.output_spec().get()
        outputs['out_file'] = os.path.abspath(self.inputs.out_file)
        return outputs


class RegisterInputSpec(FSTraitedSpec):
    # required
    in_surf = File(
        argstr="%s",
        exists=True,
        mandatory=True,
        position=-3,
        copyfile=True,
        desc="Surface to register, often {hemi}.sphere")
    target = File(
        argstr="%s",
        exists=True,
        mandatory=True,
        position=-2,
        desc="The data to register to. In normal recon-all usage, " +
        "this is a template file for average surface.")
    in_sulc = File(
        exists=True,
        mandatory=True,
        copyfile=True,
        desc=
        "Undocumented mandatory input file ${SUBJECTS_DIR}/surf/{hemisphere}.sulc "
    )
    out_file = File(
        argstr="%s",
        exists=False,
        position=-1,
        genfile=True,
        desc="Output surface file to capture registration")
    # optional
    curv = traits.Bool(
        argstr="-curv",
        requires=['in_smoothwm'],
        desc="Use smoothwm curvature for final alignment")
    in_smoothwm = File(
        exists=True,
        copyfile=True,
        desc=
        "Undocumented input file ${SUBJECTS_DIR}/surf/{hemisphere}.smoothwm ")


class RegisterOutputSpec(TraitedSpec):
    out_file = File(
        exists=False, desc="Output surface file to capture registration")


class Register(FSCommand):
    """ This program registers a surface to an average surface template.

    Examples
    ========
    >>> from nipype.interfaces.freesurfer import Register
    >>> register = Register()
    >>> register.inputs.in_surf = 'lh.pial'
    >>> register.inputs.in_smoothwm = 'lh.pial'
    >>> register.inputs.in_sulc = 'lh.pial'
    >>> register.inputs.target = 'aseg.mgz'
    >>> register.inputs.out_file = 'lh.pial.reg'
    >>> register.inputs.curv = True
    >>> register.cmdline
    'mris_register -curv lh.pial aseg.mgz lh.pial.reg'
    """

    _cmd = 'mris_register'
    input_spec = RegisterInputSpec
    output_spec = RegisterOutputSpec

    def _format_arg(self, opt, spec, val):
        if opt == 'curv':
            return spec.argstr
        return super(Register, self)._format_arg(opt, spec, val)

    def _gen_filename(self, name):
        if name == 'out_file':
            return self._list_outputs()[name]
        return None

    def _list_outputs(self):
        outputs = self.output_spec().get()
        if isdefined(self.inputs.out_file):
            outputs['out_file'] = os.path.abspath(self.inputs.out_file)
        else:
            outputs['out_file'] = os.path.abspath(self.inputs.in_surf) + '.reg'
        return outputs


class PaintInputSpec(FSTraitedSpec):
    # required
    in_surf = File(
        argstr="%s",
        exists=True,
        mandatory=True,
        position=-2,
        desc="Surface file with grid (vertices) onto which the " +
        "template data is to be sampled or 'painted'")
    template = File(
        argstr="%s",
        exists=True,
        mandatory=True,
        position=-3,
        desc="Template file")
    # optional
    template_param = traits.Int(desc="Frame number of the input template")
    averages = traits.Int(argstr="-a %d", desc="Average curvature patterns")
    out_file = File(
        argstr="%s",
        exists=False,
        position=-1,
        name_template="%s.avg_curv",
        hash_files=False,
        name_source=['in_surf'],
        keep_extension=False,
        desc="File containing a surface-worth of per-vertex values, " +
        "saved in 'curvature' format.")


class PaintOutputSpec(TraitedSpec):
    out_file = File(
        exists=False,
        desc=
        "File containing a surface-worth of per-vertex values, saved in 'curvature' format."
    )


class Paint(FSCommand):
    """
    This program is useful for extracting one of the arrays ("a variable")
    from a surface-registration template file. The output is a file
    containing a surface-worth of per-vertex values, saved in "curvature"
    format. Because the template data is sampled to a particular surface
    mesh, this conjures the idea of "painting to a surface".

    Examples
    ========
    >>> from nipype.interfaces.freesurfer import Paint
    >>> paint = Paint()
    >>> paint.inputs.in_surf = 'lh.pial'
    >>> paint.inputs.template = 'aseg.mgz'
    >>> paint.inputs.averages = 5
    >>> paint.inputs.out_file = 'lh.avg_curv'
    >>> paint.cmdline
    'mrisp_paint -a 5 aseg.mgz lh.pial lh.avg_curv'
    """

    _cmd = 'mrisp_paint'
    input_spec = PaintInputSpec
    output_spec = PaintOutputSpec

    def _format_arg(self, opt, spec, val):
        if opt == 'template':
            if isdefined(self.inputs.template_param):
                return spec.argstr % (
                    val + '#' + str(self.inputs.template_param))
        return super(Paint, self)._format_arg(opt, spec, val)

    def _list_outputs(self):
        outputs = self.output_spec().get()
        outputs['out_file'] = os.path.abspath(self.inputs.out_file)
        return outputs


class MRICoregInputSpec(FSTraitedSpec):
    source_file = File(
        argstr='--mov %s',
        desc='source file to be registered',
        mandatory=True,
        copyfile=False)
    reference_file = File(
        argstr='--ref %s',
        desc='reference (target) file',
        mandatory=True,
        copyfile=False,
        xor=['subject_id'])
    out_lta_file = traits.Either(
        True,
        File,
        argstr='--lta %s',
        default=True,
        usedefault=True,
        desc='output registration file (LTA format)')
    out_reg_file = traits.Either(
        True,
        File,
        argstr='--regdat %s',
        desc='output registration file (REG format)')
    out_params_file = traits.Either(
        True, File, argstr='--params %s', desc='output parameters file')

    subjects_dir = Directory(
        exists=True, argstr='--sd %s', desc='FreeSurfer SUBJECTS_DIR')
    subject_id = traits.Str(
        argstr='--s %s',
        position=1,
        mandatory=True,
        xor=['reference_file'],
        requires=['subjects_dir'],
        desc='freesurfer subject ID (implies ``reference_mask == '
        'aparc+aseg.mgz`` unless otherwise specified)')
    dof = traits.Enum(
        6,
        9,
        12,
        argstr='--dof %d',
        desc='number of transform degrees of freedom')
    reference_mask = traits.Either(
        False,
        traits.Str,
        argstr='--ref-mask %s',
        position=2,
        desc='mask reference volume with given mask, or None if ``False``')
    source_mask = traits.Str(
        argstr='--mov-mask', desc='mask source file with given mask')
    num_threads = traits.Int(
        argstr='--threads %d', desc='number of OpenMP threads')
    no_coord_dithering = traits.Bool(
        argstr='--no-coord-dither', desc='turn off coordinate dithering')
    no_intensity_dithering = traits.Bool(
        argstr='--no-intensity-dither', desc='turn off intensity dithering')
    sep = traits.List(
        argstr='--sep %s...',
        minlen=1,
        maxlen=2,
        desc='set spatial scales, in voxels (default [2, 4])')
    initial_translation = traits.Tuple(
        traits.Float,
        traits.Float,
        traits.Float,
        argstr='--trans %g %g %g',
        desc='initial translation in mm (implies no_cras0)')
    initial_rotation = traits.Tuple(
        traits.Float,
        traits.Float,
        traits.Float,
        argstr='--rot %g %g %g',
        desc='initial rotation in degrees')
    initial_scale = traits.Tuple(
        traits.Float,
        traits.Float,
        traits.Float,
        argstr='--scale %g %g %g',
        desc='initial scale')
    initial_shear = traits.Tuple(
        traits.Float,
        traits.Float,
        traits.Float,
        argstr='--shear %g %g %g',
        desc='initial shear (Hxy, Hxz, Hyz)')
    no_cras0 = traits.Bool(
        argstr='--no-cras0',
        desc='do not set translation parameters to align '
        'centers of source and reference files')
    max_iters = traits.Range(
        low=1, argstr='--nitersmax %d', desc='maximum iterations (default: 4)')
    ftol = traits.Float(
        argstr='--ftol %e', desc='floating-point tolerance (default=1e-7)')
    linmintol = traits.Float(argstr='--linmintol %e')
    saturation_threshold = traits.Range(
        low=0.0,
        high=100.0,
        argstr='--sat %g',
        desc='saturation threshold (default=9.999)')
    conform_reference = traits.Bool(
        argstr='--conf-ref', desc='conform reference without rescaling')
    no_brute_force = traits.Bool(
        argstr='--no-bf', desc='do not brute force search')
    brute_force_limit = traits.Float(
        argstr='--bf-lim %g',
        xor=['no_brute_force'],
        desc='constrain brute force search to +/- lim')
    brute_force_samples = traits.Int(
        argstr='--bf-nsamp %d',
        xor=['no_brute_force'],
        desc='number of samples in brute force search')
    no_smooth = traits.Bool(
        argstr='--no-smooth',
        desc='do not apply smoothing to either reference or source file')
    ref_fwhm = traits.Float(
        argstr='--ref-fwhm', desc='apply smoothing to reference file')
    source_oob = traits.Bool(
        argstr='--mov-oob',
        desc='count source voxels that are out-of-bounds as 0')
    # Skipping mat2par


class MRICoregOutputSpec(TraitedSpec):
    out_reg_file = File(exists=True, desc='output registration file')
    out_lta_file = File(exists=True, desc='output LTA-style registration file')
    out_params_file = File(exists=True, desc='output parameters file')


class MRICoreg(FSCommand):
    """ This program registers one volume to another

    mri_coreg is a C reimplementation of spm_coreg in FreeSurfer

    Examples
    ========
    >>> from nipype.interfaces.freesurfer import MRICoreg
    >>> coreg = MRICoreg()
    >>> coreg.inputs.source_file = 'moving1.nii'
    >>> coreg.inputs.reference_file = 'fixed1.nii'
    >>> coreg.inputs.subjects_dir = '.'
    >>> coreg.cmdline # doctest: +ELLIPSIS
    'mri_coreg --lta .../registration.lta --ref fixed1.nii --mov moving1.nii --sd .'

    If passing a subject ID, the reference mask may be disabled:

    >>> coreg = MRICoreg()
    >>> coreg.inputs.source_file = 'moving1.nii'
    >>> coreg.inputs.subjects_dir = '.'
    >>> coreg.inputs.subject_id = 'fsaverage'
    >>> coreg.inputs.reference_mask = False
    >>> coreg.cmdline # doctest: +ELLIPSIS
    'mri_coreg --s fsaverage --no-ref-mask --lta .../registration.lta --mov moving1.nii --sd .'

    Spatial scales may be specified as a list of one or two separations:

    >>> coreg.inputs.sep = [4]
    >>> coreg.cmdline # doctest: +ELLIPSIS
    'mri_coreg --s fsaverage --no-ref-mask --lta .../registration.lta --sep 4 --mov moving1.nii --sd .'

    >>> coreg.inputs.sep = [4, 5]
    >>> coreg.cmdline # doctest: +ELLIPSIS
    'mri_coreg --s fsaverage --no-ref-mask --lta .../registration.lta --sep 4 --sep 5 --mov moving1.nii --sd .'
    """

    _cmd = 'mri_coreg'
    input_spec = MRICoregInputSpec
    output_spec = MRICoregOutputSpec

    def _format_arg(self, opt, spec, val):
        if opt in ('out_reg_file', 'out_lta_file',
                   'out_params_file') and val is True:
            val = self._list_outputs()[opt]
        elif opt == 'reference_mask' and val is False:
            return '--no-ref-mask'
        return super(MRICoreg, self)._format_arg(opt, spec, val)

    def _list_outputs(self):
        outputs = self.output_spec().get()

        out_lta_file = self.inputs.out_lta_file
        if isdefined(out_lta_file):
            if out_lta_file is True:
                out_lta_file = 'registration.lta'
            outputs['out_lta_file'] = os.path.abspath(out_lta_file)

        out_reg_file = self.inputs.out_reg_file
        if isdefined(out_reg_file):
            if out_reg_file is True:
                out_reg_file = 'registration.dat'
            outputs['out_reg_file'] = os.path.abspath(out_reg_file)

        out_params_file = self.inputs.out_params_file
        if isdefined(out_params_file):
            if out_params_file is True:
                out_params_file = 'registration.par'
            outputs['out_params_file'] = os.path.abspath(out_params_file)

        return outputs
