# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
"""
The QT1 module of niftyfit, which wraps the Multi-Echo T1 fitting methods
in NiftyFit.
"""

from ..base import TraitedSpec, File, traits, CommandLineInputSpec
from .base import NiftyFitCommand
from ..niftyreg.base import get_custom_path


class FitQt1InputSpec(CommandLineInputSpec):
    """ Input Spec for FitQt1. """
    desc = 'Filename of the 4D Multi-Echo T1 source image.'
    source_file = File(
        position=1,
        exists=True,
        desc=desc,
        argstr='-source %s',
        mandatory=True)

    # Output options:
    t1map_file = File(
        name_source=['source_file'],
        name_template='%s_t1map.nii.gz',
        argstr='-t1map %s',
        desc='Filename of the estimated output T1 map (in ms).')
    m0map_file = File(
        name_source=['source_file'],
        name_template='%s_m0map.nii.gz',
        argstr='-m0map %s',
        desc='Filename of the estimated input M0 map.')
    desc = 'Filename of the estimated output multi-parameter map.'
    mcmap_file = File(
        name_source=['source_file'],
        name_template='%s_mcmap.nii.gz',
        argstr='-mcmap %s',
        desc=desc)
    comp_file = File(
        name_source=['source_file'],
        name_template='%s_comp.nii.gz',
        argstr='-comp %s',
        desc='Filename of the estimated multi-component T1 map.')
    desc = 'Filename of the error map (symmetric matrix, [Diag,OffDiag]).'
    error_file = File(
        name_source=['source_file'],
        name_template='%s_error.nii.gz',
        argstr='-error %s',
        desc=desc)
    syn_file = File(
        name_source=['source_file'],
        name_template='%s_syn.nii.gz',
        argstr='-syn %s',
        desc='Filename of the synthetic ASL data.')
    res_file = File(
        name_source=['source_file'],
        name_template='%s_res.nii.gz',
        argstr='-res %s',
        desc='Filename of the model fit residuals')

    # Other options:
    mask = File(
        position=2,
        exists=True,
        desc='Filename of image mask.',
        argstr='-mask %s')
    prior = File(
        position=3,
        exists=True,
        desc='Filename of parameter prior.',
        argstr='-prior %s')
    te_value = traits.Float(
        desc='TE Echo Time [0ms!].', argstr='-TE %f', position=4)
    tr_value = traits.Float(
        desc='TR Repetition Time [10s!].', argstr='-TR %f', position=5)
    desc = 'Number of components to fit [1] (currently IR/SR only)'
    # set position to be ahead of TIs
    nb_comp = traits.Int(desc=desc, position=6, argstr='-nc %d')
    desc = 'Set LM parameters (initial value, decrease rate) [100,1.2].'
    lm_val = traits.Tuple(
        traits.Float, traits.Float, desc=desc, argstr='-lm %f %f', position=7)
    desc = 'Use Gauss-Newton algorithm [Levenberg-Marquardt].'
    gn_flag = traits.Bool(desc=desc, argstr='-gn', position=8)
    slice_no = traits.Int(
        desc='Fit to single slice number.', argstr='-slice %d', position=9)
    voxel = traits.Tuple(
        traits.Int,
        traits.Int,
        traits.Int,
        desc='Fit to single voxel only.',
        argstr='-voxel %d %d %d',
        position=10)
    maxit = traits.Int(
        desc='NLSQR iterations [100].', argstr='-maxit %d', position=11)

    # IR options:
    sr_flag = traits.Bool(
        desc='Saturation Recovery fitting [default].',
        argstr='-SR',
        position=12)
    ir_flag = traits.Bool(
        desc='Inversion Recovery fitting [default].',
        argstr='-IR',
        position=13)
    tis = traits.List(
        traits.Float,
        position=14,
        desc='Inversion times for T1 data [1s,2s,5s].',
        argstr='-TIs %s',
        sep=' ')
    tis_list = File(
        exists=True,
        argstr='-TIlist %s',
        desc='Filename of list of pre-defined TIs.')
    t1_list = File(
        exists=True,
        argstr='-T1list %s',
        desc='Filename of list of pre-defined T1s')
    t1min = traits.Float(
        desc='Minimum tissue T1 value [400ms].', argstr='-T1min %f')
    t1max = traits.Float(
        desc='Maximum tissue T1 value [4000ms].', argstr='-T1max %f')

    # SPGR options
    spgr = traits.Bool(desc='Spoiled Gradient Echo fitting', argstr='-SPGR')
    flips = traits.List(
        traits.Float, desc='Flip angles', argstr='-flips %s', sep=' ')
    desc = 'Filename of list of pre-defined flip angles (deg).'
    flips_list = File(exists=True, argstr='-fliplist %s', desc=desc)
    desc = 'Filename of B1 estimate for fitting (or include in prior).'
    b1map = File(exists=True, argstr='-b1map %s', desc=desc)

    # MCMC options:
    mcout = File(
        exists=True,
        desc='Filename of mc samples (ascii text file)',
        argstr='-mcout %s')
    mcsamples = traits.Int(
        desc='Number of samples to keep [100].', argstr='-mcsamples %d')
    mcmaxit = traits.Int(
        desc='Number of iterations to run [10,000].', argstr='-mcmaxit %d')
    acceptance = traits.Float(
        desc='Fraction of iterations to accept [0.23].',
        argstr='-acceptance %f')


class FitQt1OutputSpec(TraitedSpec):
    """ Output Spec for FitQt1. """
    t1map_file = File(desc='Filename of the estimated output T1 map (in ms)')
    m0map_file = File(desc='Filename of the m0 map')
    desc = 'Filename of the estimated output multi-parameter map'
    mcmap_file = File(desc=desc)
    comp_file = File(desc='Filename of the estimated multi-component T1 map.')
    desc = 'Filename of the error map (symmetric matrix, [Diag,OffDiag])'
    error_file = File(desc=desc)
    syn_file = File(desc='Filename of the synthetic ASL data')
    res_file = File(desc='Filename of the model fit residuals')


class FitQt1(NiftyFitCommand):
    """Interface for executable fit_qt1 from Niftyfit platform.

    Use NiftyFit to perform Qt1 fitting.

    T1 Fitting Routine (To inversion recovery or spgr data).
    Fits single component T1 maps in the first instance.

    `Source code <https://cmiclab.cs.ucl.ac.uk/CMIC/NiftyFit-Release>`_

    Examples
    --------

    >>> from nipype.interfaces.niftyfit import FitQt1
    >>> fit_qt1 = FitQt1()
    >>> fit_qt1.inputs.source_file = 'TI4D.nii.gz'
    >>> fit_qt1.cmdline
    'fit_qt1 -source TI4D.nii.gz -comp TI4D_comp.nii.gz \
-error TI4D_error.nii.gz -m0map TI4D_m0map.nii.gz -mcmap TI4D_mcmap.nii.gz \
-res TI4D_res.nii.gz -syn TI4D_syn.nii.gz -t1map TI4D_t1map.nii.gz'

    """
    _cmd = get_custom_path('fit_qt1', env_dir='NIFTYFITDIR')
    input_spec = FitQt1InputSpec
    output_spec = FitQt1OutputSpec
    _suffix = '_fit_qt1'
