#!/usr/bin/env python
#ICA_AROMA pulled from: https://github.com/rhr-pruim/ICA-AROMA

from nipype.interfaces.base import (
    TraitedSpec,
    CommandLineInputSpec,
    CommandLine,
    File,
    Directory,
    traits
)
import os

class ICA_AROMAInputSpec(CommandLineInputSpec):
    #How do I format different processing streams (featdir versus manual inputs)?
    #e.g. I want featDir to be mutually exclusive with infile,mask,affmat,warp, and mc.
    #....xor can make all the options mutually exclusive with eachother, but I want
    #to group the options (infile,mask,affmat,warp,mc)
    #idk how xor works, my attempt
    _xor_inputs = ('featDir','infile')
    featDir = Directory(exists=True,
                        desc='If a feat directory exists and temporal filtering '
                        'has not been run yet, ICA_AROMA can use the files in '
                        'this directory.',mandatory=False,xor=_xor_inputs)
    infile = File(exists=True,
                  desc='volume to be denoised',
                  argstr='-i %s',xor=_xor_inputs)
    outdir = Directory(desc='path to output directory',
                  argstr='-o %s',mandatory=True)
    mask = File(exists=True,
                desc='path/name volume mask',
                argstr='-m %s')
    dim = traits.Int(desc='Dimensionality reduction when running '
                     'MELODIC (defualt is automatic estimation)',
                     argstr='-dim %d')
    TR = traits.Float(desc='TR in seconds. If this is not specified '
                      'the TR will be extracted from the header of the fMRI nifti file.',
                      argstr='%.2f')
    melodic_dir = Directory(exists=True,
                            desc='path to MELODIC directory if MELODIC has already been ran',
                            argstr='-meldir %s')
    affmat = File(exists=True,
                  desc='path/name of the mat-file describing the '
                  'affine registration (e.g. FSL FLIRT) of the '
                  'functional data to structural space (.mat file)',
                  argstr='-affmat %s')
    warp = File(exists=True,
                desc='File name of the warp-file describing '
                'the non-linear registration (e.g. FSL FNIRT) '
                'of the structural data to MNI152 space (.nii.gz)',
                argstr='-warp %s')
    mc = File(exists=True,
              desc='motion parameters file',
              argstr='-mc %s')
    denoise_type = traits.Str(argstr='-den %s',
                              desc='Type of denoising strategy: '
                              '-none: only classification, no denoising '
                              '-nonaggr (default): non-aggresssive denoising, i.e. partial component regression '
                              '-aggr: aggressive denoising, i.e. full component regression '
                              '-both: both aggressive and non-aggressive denoising (two outputs)',
                              mandatory=True)

class ICA_AROMAOutputSpec(TraitedSpec):
    denoised_func_data_nonaggr = File(desc='if generated: denoised volume (non aggressive)')
    denoised_func_data_aggr = File(desc='if generated: denoised volume (aggressive)')

class ICA_AROMA(CommandLine):
    _cmd = 'ICA_AROMA.py'
    input_spec = ICA_AROMAInputSpec
    output_spec = ICA_AROMAOutputSpec

    def _list_outputs(self):
        #How do the outputs output_spec actually relate to the output filenames?
        outputs = self.output_spec.get()
        outdir = self.input_spec.outdir
        denoising_strategy = input_spec.denoise_type

        if denoising_strategy is "noaggr":
          out_file = os.path.join(outdir,'denoised_func_data_nonaggr.nii.gz')
        elif denoising_strategy is "aggr":
          out_file = os.path.join(outdir,'denoised_func_data_aggr.nii.gz')
        elif denoising_strategy is "both":
          #this is a tuple, versus the single outputs, how do I return a variable number of outputs?
          out_file = (os.path.join(outdir,'denoised_func_data_nonaggr.nii.gz'), os.path.join(outdir,'denoised_func_data_aggr.nii.gz'))
        elif denoising_strategy is "none":
          print "No denoising selected"
        else:
          raise RuntimeError('denoise_type must be specified as one of'
                             ' noaggr,aggr,both, or none')

        return out_file