# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
# -*- coding: utf-8 -*-
from __future__ import (print_function, division, unicode_literals,
                        absolute_import)

import os
import os.path as op

from ..base import (CommandLineInputSpec, traits, TraitedSpec, File, isdefined)
from .base import MRTrix3Base


class BuildConnectomeInputSpec(CommandLineInputSpec):
    in_file = File(
        exists=True,
        argstr='%s',
        mandatory=True,
        position=-3,
        desc='input tractography')
    in_parc = File(
        exists=True, argstr='%s', position=-2, desc='parcellation file')
    out_file = File(
        'connectome.csv',
        argstr='%s',
        mandatory=True,
        position=-1,
        usedefault=True,
        desc='output file after processing')

    nthreads = traits.Int(
        argstr='-nthreads %d',
        desc='number of threads. if zero, the number'
        ' of available cpus will be used',
        nohash=True)

    vox_lookup = traits.Bool(
        argstr='-assignment_voxel_lookup',
        desc='use a simple voxel lookup value at each streamline endpoint')
    search_radius = traits.Float(
        argstr='-assignment_radial_search %f',
        desc='perform a radial search from each streamline endpoint to locate '
        'the nearest node. Argument is the maximum radius in mm; if no node is'
        ' found within this radius, the streamline endpoint is not assigned to'
        ' any node.')
    search_reverse = traits.Float(
        argstr='-assignment_reverse_search %f',
        desc='traverse from each streamline endpoint inwards along the '
        'streamline, in search of the last node traversed by the streamline. '
        'Argument is the maximum traversal length in mm (set to 0 to allow '
        'search to continue to the streamline midpoint).')
    search_forward = traits.Float(
        argstr='-assignment_forward_search %f',
        desc='project the streamline forwards from the endpoint in search of a'
        'parcellation node voxel. Argument is the maximum traversal length in '
        'mm.')

    metric = traits.Enum(
        'count',
        'meanlength',
        'invlength',
        'invnodevolume',
        'mean_scalar',
        'invlength_invnodevolume',
        argstr='-metric %s',
        desc='specify the edge'
        ' weight metric')

    in_scalar = File(
        exists=True,
        argstr='-image %s',
        desc='provide the associated image '
        'for the mean_scalar metric')

    in_weights = File(
        exists=True,
        argstr='-tck_weights_in %s',
        desc='specify a text scalar '
        'file containing the streamline weights')

    keep_unassigned = traits.Bool(
        argstr='-keep_unassigned',
        desc='By default, the program discards the'
        ' information regarding those streamlines that are not successfully '
        'assigned to a node pair. Set this option to keep these values (will '
        'be the first row/column in the output matrix)')
    zero_diagonal = traits.Bool(
        argstr='-zero_diagonal',
        desc='set all diagonal entries in the matrix '
        'to zero (these represent streamlines that connect to the same node at'
        ' both ends)')


class BuildConnectomeOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc='the output response file')


class BuildConnectome(MRTrix3Base):
    """
    Generate a connectome matrix from a streamlines file and a node
    parcellation image

    Example
    -------

    >>> import nipype.interfaces.mrtrix3 as mrt
    >>> mat = mrt.BuildConnectome()
    >>> mat.inputs.in_file = 'tracks.tck'
    >>> mat.inputs.in_parc = 'aparc+aseg.nii'
    >>> mat.cmdline                               # doctest: +ELLIPSIS
    'tck2connectome tracks.tck aparc+aseg.nii connectome.csv'
    >>> mat.run()                                 # doctest: +SKIP
    """

    _cmd = 'tck2connectome'
    input_spec = BuildConnectomeInputSpec
    output_spec = BuildConnectomeOutputSpec

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


class LabelConfigInputSpec(CommandLineInputSpec):
    in_file = File(
        exists=True,
        argstr='%s',
        mandatory=True,
        position=-3,
        desc='input anatomical image')
    in_config = File(
        exists=True,
        argstr='%s',
        position=-2,
        desc='connectome configuration file')
    out_file = File(
        'parcellation.mif',
        argstr='%s',
        mandatory=True,
        position=-1,
        usedefault=True,
        desc='output file after processing')

    lut_basic = File(
        argstr='-lut_basic %s',
        desc='get information from '
        'a basic lookup table consisting of index / name pairs')
    lut_fs = File(
        argstr='-lut_freesurfer %s',
        desc='get information from '
        'a FreeSurfer lookup table(typically "FreeSurferColorLUT'
        '.txt")')
    lut_aal = File(
        argstr='-lut_aal %s',
        desc='get information from the AAL '
        'lookup table (typically "ROI_MNI_V4.txt")')
    lut_itksnap = File(
        argstr='-lut_itksnap %s',
        desc='get information from an'
        ' ITK - SNAP lookup table(this includes the IIT atlas '
        'file "LUT_GM.txt")')
    spine = File(
        argstr='-spine %s',
        desc='provide a manually-defined '
        'segmentation of the base of the spine where the streamlines'
        ' terminate, so that this can become a node in the connection'
        ' matrix.')
    nthreads = traits.Int(
        argstr='-nthreads %d',
        desc='number of threads. if zero, the number'
        ' of available cpus will be used',
        nohash=True)


class LabelConfigOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc='the output response file')


class LabelConfig(MRTrix3Base):
    """
    Re-configure parcellation to be incrementally defined.

    Example
    -------

    >>> import nipype.interfaces.mrtrix3 as mrt
    >>> labels = mrt.LabelConfig()
    >>> labels.inputs.in_file = 'aparc+aseg.nii'
    >>> labels.inputs.in_config = 'mrtrix3_labelconfig.txt'
    >>> labels.cmdline                               # doctest: +ELLIPSIS
    'labelconfig aparc+aseg.nii mrtrix3_labelconfig.txt parcellation.mif'
    >>> labels.run()                                 # doctest: +SKIP
    """

    _cmd = 'labelconfig'
    input_spec = LabelConfigInputSpec
    output_spec = LabelConfigOutputSpec

    def _parse_inputs(self, skip=None):
        if skip is None:
            skip = []

        if not isdefined(self.inputs.in_config):
            from distutils.spawn import find_executable
            path = find_executable(self._cmd)
            if path is None:
                path = os.getenv(MRTRIX3_HOME, '/opt/mrtrix3')
            else:
                path = op.dirname(op.dirname(path))

            self.inputs.in_config = op.join(
                path, 'src/dwi/tractography/connectomics/'
                'example_configs/fs_default.txt')

        return super(LabelConfig, self)._parse_inputs(skip=skip)

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


class LabelConvertInputSpec(CommandLineInputSpec):
    in_file = File(
        exists=True,
        argstr='%s',
        mandatory=True,
        position=-4,
        desc='input anatomical image')
    in_lut = File(
        exists=True,
        argstr='%s',
        mandatory=True,
        position=-3,
        desc='get information from '
        'a basic lookup table consisting of index / name pairs')
    in_config = File(
        exists=True,
        argstr='%s',
        position=-2,
        desc='connectome configuration file')
    out_file = File(
        'parcellation.mif',
        argstr='%s',
        mandatory=True,
        position=-1,
        usedefault=True,
        desc='output file after processing')
    spine = File(
        argstr='-spine %s',
        desc='provide a manually-defined '
        'segmentation of the base of the spine where the streamlines'
        ' terminate, so that this can become a node in the connection'
        ' matrix.')
    num_threads = traits.Int(
        argstr='-nthreads %d',
        desc='number of threads. if zero, the number'
        ' of available cpus will be used',
        nohash=True)


class LabelConvertOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc='the output response file')


class LabelConvert(MRTrix3Base):
    """
    Re-configure parcellation to be incrementally defined.

    Example
    -------

    >>> import nipype.interfaces.mrtrix3 as mrt
    >>> labels = mrt.LabelConvert()
    >>> labels.inputs.in_file = 'aparc+aseg.nii'
    >>> labels.inputs.in_config = 'mrtrix3_labelconfig.txt'
    >>> labels.inputs.in_lut = 'FreeSurferColorLUT.txt'
    >>> labels.cmdline
    'labelconvert aparc+aseg.nii FreeSurferColorLUT.txt mrtrix3_labelconfig.txt parcellation.mif'
    >>> labels.run()                                 # doctest: +SKIP
    """

    _cmd = 'labelconvert'
    input_spec = LabelConvertInputSpec
    output_spec = LabelConvertOutputSpec

    def _parse_inputs(self, skip=None):
        if skip is None:
            skip = []

        if not isdefined(self.inputs.in_config):
            from nipype.utils.filemanip import which
            path = which(self._cmd)
            if path is None:
                path = os.getenv(MRTRIX3_HOME, '/opt/mrtrix3')
            else:
                path = op.dirname(op.dirname(path))

            self.inputs.in_config = op.join(
                path, 'src/dwi/tractography/connectomics/'
                'example_configs/fs_default.txt')

        return super(LabelConvert, self)._parse_inputs(skip=skip)

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