# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:

import os
from tempfile import mkdtemp
from shutil import rmtree

import numpy as np

import nibabel as nb
from nipype.testing import (assert_equal, assert_not_equal,
                            assert_raises, skipif)
import nipype.interfaces.fsl.utils as fsl
from nipype.interfaces.fsl import no_fsl


def create_files_in_directory():
    outdir = mkdtemp()
    cwd = os.getcwd()
    os.chdir(outdir)
    filelist = ['a.nii', 'b.nii']
    for f in filelist:
        hdr = nb.Nifti1Header()
        shape = (3, 3, 3, 4)
        hdr.set_data_shape(shape)
        img = np.random.random(shape)
        nb.save(nb.Nifti1Image(img, np.eye(4), hdr),
                os.path.join(outdir, f))
    return filelist, outdir, cwd


def clean_directory(outdir, old_wd):
    if os.path.exists(outdir):
        rmtree(outdir)
    os.chdir(old_wd)


@skipif(no_fsl)
def test_extractroi():
    input_map = dict(args=dict(argstr='%s',),
                     environ=dict(),
                     in_file=dict(mandatory=True, argstr='%s',),
                     output_type=dict(),
                     roi_file=dict(argstr='%s',),
                     t_min=dict(argstr='%d',),
                     t_size=dict(argstr='%d',),
                     x_min=dict(argstr='%d',),
                     x_size=dict(argstr='%d',),
                     y_min=dict(argstr='%d',),
                     y_size=dict(argstr='%d',),
                     z_min=dict(argstr='%d',),
                     z_size=dict(argstr='%d',),
                     )
    instance = fsl.ExtractROI()
    for key, metadata in input_map.items():
        for metakey, value in metadata.items():
            yield assert_equal, getattr(instance.inputs.traits()[key], metakey), value


@skipif(no_fsl)
def test_imagemaths():
    input_map = dict(args=dict(argstr='%s',),
                     environ=dict(),
                     in_file=dict(argstr='%s', mandatory=True,),
                     in_file2=dict(argstr='%s',),
                     op_string=dict(argstr='%s',),
                     out_data_type=dict(argstr='-odt %s',),
                     out_file=dict(argstr='%s',),
                     output_type=dict(),
                     suffix=dict(),
                     )
    instance = fsl.ImageMaths()
    for key, metadata in input_map.items():
        for metakey, value in metadata.items():
            yield assert_equal, getattr(instance.inputs.traits()[key], metakey), value


@skipif(no_fsl)
def test_merge():
    input_map = dict(args=dict(argstr='%s',),
                     dimension=dict(argstr='-%s', mandatory=True,),
                     environ=dict(),
                     in_files=dict(mandatory=True, argstr='%s',),
                     merged_file=dict(argstr='%s',),
                     output_type=dict(),
                     )
    instance = fsl.Merge()
    for key, metadata in input_map.items():
        for metakey, value in metadata.items():
            yield assert_equal, getattr(instance.inputs.traits()[key], metakey), value


@skipif(no_fsl)
def test_filterregressor():
    input_map = dict(Out_vnscales=dict(),
                     args=dict(argstr='%s',),
                     design_file=dict(mandatory=True,),
                     environ=dict(),
                     filter_columns=dict(mandatory=True,),
                     filter_all=dict(mandatory=True,),
                     in_file=dict(mandatory=True,),
                     mask=dict(),
                     out_file=dict(),
                     output_type=dict(),
                     var_norm=dict(),
                     )
    instance = fsl.FilterRegressor()
    for key, metadata in input_map.items():
        for metakey, value in metadata.items():
            yield assert_equal, getattr(instance.inputs.traits()[key], metakey), value


@skipif(no_fsl)
def test_smooth():
    input_map = dict(args=dict(argstr='%s',),
                     environ=dict(),
                     fwhm=dict(
                         argstr='-kernel gauss %f -fmean', mandatory=True,),
                     in_file=dict(argstr='%s', mandatory=True,),
                     output_type=dict(),
                     smoothed_file=dict(argstr='%s',),
                     )
    instance = fsl.Smooth()
    for key, metadata in input_map.items():
        for metakey, value in metadata.items():
            yield assert_equal, getattr(instance.inputs.traits()[key], metakey), value


@skipif(no_fsl)
def test_split():
    input_map = dict(args=dict(argstr='%s',),
                     dimension=dict(argstr='-%s',),
                     environ=dict(),
                     in_file=dict(argstr='%s',),
                     out_base_name=dict(argstr='%s',),
                     output_type=dict(),
                     )
    instance = fsl.Split()
    for key, metadata in input_map.items():
        for metakey, value in metadata.items():
            yield assert_equal, getattr(instance.inputs.traits()[key], metakey), value


def no_fsl():
    """Checks if FSL is NOT installed
    used with skipif to skip tests that will
    fail if FSL is not installed"""

    if fsl.Info().version() == None:
        return True
    else:
        return False


@skipif(no_fsl)
def test_fslroi():
    filelist, outdir, cwd = create_files_in_directory()

    roi = fsl.ExtractROI()

    # make sure command gets called
    yield assert_equal, roi.cmd, 'fslroi'

    # test raising error with mandatory args absent
    yield assert_raises, ValueError, roi.run

    # .inputs based parameters setting
    roi.inputs.in_file = filelist[0]
    roi.inputs.roi_file = 'foo_roi.nii'
    roi.inputs.t_min = 10
    roi.inputs.t_size = 20
    yield assert_equal, roi.cmdline, 'fslroi %s foo_roi.nii 10 20' % filelist[0]

    # .run based parameter setting
    roi2 = fsl.ExtractROI(in_file=filelist[0],
                          roi_file='foo2_roi.nii',
                          t_min=20, t_size=40,
                          x_min=3, x_size=30,
                          y_min=40, y_size=10,
                          z_min=5, z_size=20)
    yield assert_equal, roi2.cmdline, \
        'fslroi %s foo2_roi.nii 3 30 40 10 5 20 20 40' % filelist[0]

    clean_directory(outdir, cwd)
    # test arguments for opt_map
    # Fslroi class doesn't have a filled opt_map{}


# test fslmath
@skipif(no_fsl)
def test_fslmaths():
    filelist, outdir, cwd = create_files_in_directory()
    math = fsl.ImageMaths()

    # make sure command gets called
    yield assert_equal, math.cmd, 'fslmaths'

    # test raising error with mandatory args absent
    yield assert_raises, ValueError, math.run

    # .inputs based parameters setting
    math.inputs.in_file = filelist[0]
    math.inputs.op_string = '-add 2.5 -mul input_volume2'
    math.inputs.out_file = 'foo_math.nii'
    yield assert_equal, math.cmdline, \
        'fslmaths %s -add 2.5 -mul input_volume2 foo_math.nii' % filelist[0]

    # .run based parameter setting
    math2 = fsl.ImageMaths(in_file=filelist[0], op_string='-add 2.5',
                           out_file='foo2_math.nii')
    yield assert_equal, math2.cmdline, 'fslmaths %s -add 2.5 foo2_math.nii' % filelist[0]

    # test arguments for opt_map
    # Fslmath class doesn't have opt_map{}
    clean_directory(outdir, cwd)

# test overlay


@skipif(no_fsl)
def test_overlay():
    filelist, outdir, cwd = create_files_in_directory()
    overlay = fsl.Overlay()

    # make sure command gets called
    yield assert_equal, overlay.cmd, 'overlay'

    # test raising error with mandatory args absent
    yield assert_raises, ValueError, overlay.run

    # .inputs based parameters setting
    overlay.inputs.stat_image = filelist[0]
    overlay.inputs.stat_thresh = (2.5, 10)
    overlay.inputs.background_image = filelist[1]
    overlay.inputs.auto_thresh_bg = True
    overlay.inputs.show_negative_stats = True
    overlay.inputs.out_file = 'foo_overlay.nii'
    yield assert_equal, overlay.cmdline, \
        'overlay 1 0 %s -a %s 2.50 10.00 %s -2.50 -10.00 foo_overlay.nii' % (
            filelist[1], filelist[0], filelist[0])

    # .run based parameter setting
    overlay2 = fsl.Overlay(stat_image=filelist[0], stat_thresh=(2.5, 10),
                           background_image=filelist[1], auto_thresh_bg=True,
                           out_file='foo2_overlay.nii')
    yield assert_equal, overlay2.cmdline, 'overlay 1 0 %s -a %s 2.50 10.00 foo2_overlay.nii' % (
        filelist[1], filelist[0])

    clean_directory(outdir, cwd)

# test slicer


@skipif(no_fsl)
def test_slicer():
    filelist, outdir, cwd = create_files_in_directory()
    slicer = fsl.Slicer()

    # make sure command gets called
    yield assert_equal, slicer.cmd, 'slicer'

    # test raising error with mandatory args absent
    yield assert_raises, ValueError, slicer.run

    # .inputs based parameters setting
    slicer.inputs.in_file = filelist[0]
    slicer.inputs.image_edges = filelist[1]
    slicer.inputs.intensity_range = (10., 20.)
    slicer.inputs.all_axial = True
    slicer.inputs.image_width = 750
    slicer.inputs.out_file = 'foo_bar.png'
    yield assert_equal, slicer.cmdline, \
        'slicer %s %s -L -i 10.000 20.000  -A 750 foo_bar.png' % (
            filelist[0], filelist[1])

    # .run based parameter setting
    slicer2 = fsl.Slicer(
        in_file=filelist[0], middle_slices=True, label_slices=False,
        out_file='foo_bar2.png')
    yield assert_equal, slicer2.cmdline, 'slicer %s   -a foo_bar2.png' % (filelist[0])

    clean_directory(outdir, cwd)


def create_parfiles():

    np.savetxt('a.par', np.random.rand(6, 3))
    np.savetxt('b.par', np.random.rand(6, 3))
    return ['a.par', 'b.par']

# test fsl_tsplot


@skipif(no_fsl)
def test_plottimeseries():
    filelist, outdir, cwd = create_files_in_directory()
    parfiles = create_parfiles()
    plotter = fsl.PlotTimeSeries()

    # make sure command gets called
    yield assert_equal, plotter.cmd, 'fsl_tsplot'

    # test raising error with mandatory args absent
    yield assert_raises, ValueError, plotter.run

    # .inputs based parameters setting
    plotter.inputs.in_file = parfiles[0]
    plotter.inputs.labels = ['x', 'y', 'z']
    plotter.inputs.y_range = (0, 1)
    plotter.inputs.title = 'test plot'
    plotter.inputs.out_file = 'foo.png'
    yield assert_equal, plotter.cmdline, \
        ('fsl_tsplot -i %s -a x,y,z -o foo.png -t \'test plot\' -u 1 --ymin=0 --ymax=1'
         % parfiles[0])

    # .run based parameter setting
    plotter2 = fsl.PlotTimeSeries(
        in_file=parfiles, title='test2 plot', plot_range=(2, 5),
        out_file='bar.png')
    yield assert_equal, plotter2.cmdline, \
        'fsl_tsplot -i %s,%s -o bar.png --start=2 --finish=5 -t \'test2 plot\' -u 1' % tuple(
            parfiles)

    clean_directory(outdir, cwd)


@skipif(no_fsl)
def test_plotmotionparams():
    filelist, outdir, cwd = create_files_in_directory()
    parfiles = create_parfiles()
    plotter = fsl.PlotMotionParams()

    # make sure command gets called
    yield assert_equal, plotter.cmd, 'fsl_tsplot'

    # test raising error with mandatory args absent
    yield assert_raises, ValueError, plotter.run

    # .inputs based parameters setting
    plotter.inputs.in_file = parfiles[0]
    plotter.inputs.in_source = 'fsl'
    plotter.inputs.plot_type = 'rotations'
    plotter.inputs.out_file = 'foo.png'
    yield assert_equal, plotter.cmdline, \
        ('fsl_tsplot -i %s -o foo.png -t \'MCFLIRT estimated rotations (radians)\' '
         '--start=1 --finish=3 -a x,y,z' % parfiles[0])

    # .run based parameter setting
    plotter2 = fsl.PlotMotionParams(
        in_file=parfiles[1], in_source='spm', plot_type='translations',
        out_file='bar.png')
    yield assert_equal, plotter2.cmdline, \
        ('fsl_tsplot -i %s -o bar.png -t \'Realign estimated translations (mm)\' '
         '--start=1 --finish=3 -a x,y,z' % parfiles[1])

    clean_directory(outdir, cwd)


@skipif(no_fsl)
def test_convertxfm():
    filelist, outdir, cwd = create_files_in_directory()
    cvt = fsl.ConvertXFM()

    # make sure command gets called
    yield assert_equal, cvt.cmd, "convert_xfm"

    # test raising error with mandatory args absent
    yield assert_raises, ValueError, cvt.run

    # .inputs based parameters setting
    cvt.inputs.in_file = filelist[0]
    cvt.inputs.invert_xfm = True
    cvt.inputs.out_file = "foo.mat"
    yield assert_equal, cvt.cmdline, 'convert_xfm -omat foo.mat -inverse %s' % filelist[0]

    # constructor based parameter setting
    cvt2 = fsl.ConvertXFM(
        in_file=filelist[0], in_file2=filelist[1], concat_xfm=True,
        out_file="bar.mat")
    yield assert_equal, cvt2.cmdline, \
        "convert_xfm -omat bar.mat -concat %s %s" % (filelist[1], filelist[0])

    clean_directory(outdir, cwd)


@skipif(no_fsl)
def test_swapdims():
    files, testdir, origdir = create_files_in_directory()
    swap = fsl.SwapDimensions()

    # Test the underlying command
    yield assert_equal, swap.cmd, "fslswapdim"

    # Test mandatory args
    args = [dict(in_file=files[0]), dict(new_dims=("x", "y", "z"))]
    for arg in args:
        wontrun = fsl.SwapDimensions(**arg)
        yield assert_raises, ValueError, wontrun.run

    # Now test a basic command line
    swap.inputs.in_file = files[0]
    swap.inputs.new_dims = ("x", "y", "z")
    yield assert_equal, swap.cmdline, "fslswapdim a.nii x y z %s" % os.path.realpath(os.path.join(testdir, "a_newdims.nii"))

    # Test that we can set an output name
    swap.inputs.out_file = "b.nii"
    yield assert_equal, swap.cmdline, "fslswapdim a.nii x y z b.nii"

    # Clean up
    clean_directory(testdir, origdir)


@skipif(no_fsl)
def test_invwarp():
    input_map = dict(environ=dict(),
                     args=dict(argstr='%s',),
                     warp=dict(mandatory=True,),
                     reference=dict(mandatory=True,))
    instance = fsl.InvWarp()
    for key, metadata in input_map.items():
        for metakey, value in metadata.items():
            yield assert_equal, getattr(instance.inputs.traits()[key], metakey), value
