def load_vols(vols):
    if isinstance(vols, list):
        return [load_vol(vol) if isinstance(vol, basestring)
                else vol for vol in vols]
    elif isinstance(vols, basestring):
        return nibabel.four_to_three(nibabel.load(vols))
    elif is_niimg(vols):
        return nibabel.four_to_three(vols)
    else:
        raise TypeError(type(vols))

    def prep_data(self, nifti1, bval1, bvec1, nifti2, bval2, bvec2):
        ''' Load the reconstructed image files and generate the files that TOPUP needs. '''
        ni1 = nb.load(nifti1)
        ni2 = nb.load(nifti2)
        phase_dim1 = ni1.get_header().get_dim_info()[1]
        phase_dim2 = ni2.get_header().get_dim_info()[1]

        bvals1 = np.loadtxt(bval1)
        bvals2 = np.loadtxt(bval2)
        bvecs1 = np.loadtxt(bvec1)
        bvecs2 = np.loadtxt(bvec2)

        nondwi1 = [im for i,im in enumerate(nb.four_to_three(ni1)) if bvals1[i]<10 and i<self.num_vols]
        nondwi2 = [im for i,im in enumerate(nb.four_to_three(ni2)) if bvals2[i]<10 and i<self.num_vols]

        b0 = nb.concat_images(nondwi1+nondwi2)
        # Topup requires an even number of slices
        if b0.shape[2]%2:
            d = b0.get_data()
            d = np.concatenate((d,np.zeros((d.shape[0],d.shape[1],1,d.shape[3]), dtype=d.dtype)),axis=2)
            b0 = nb.Nifti1Image(d, b0.get_affine())

        nb.save(b0, self.b0_file)
        with open(self.acq_file, 'w') as f:
            for i in xrange(len(nondwi1)):
                row = ['0','0','0',str(self.readout_time1),'\n']
                row[phase_dim1] = str(self.pe_dir1)
                f.write(' '.join(row))
            for i in xrange(len(nondwi2)):
                row = ['0','0','0',str(self.readout_time2),'\n']
                row[phase_dim2] = str(self.pe_dir2)
                f.write(' '.join(row))

        mux_ims1 = nb.four_to_three(ni1)[self.num_cal1:]
        mux_ims2 = nb.four_to_three(ni2)[self.num_cal2:]
        all_ims = nb.concat_images(mux_ims1 + mux_ims2)
        if all_ims.shape[2]%2:
            d = all_ims.get_data()
            d = np.concatenate((d,np.zeros((d.shape[0],d.shape[1],1,d.shape[3]), dtype=d.dtype)),axis=2)
            all_ims = nb.Nifti1Image(d, all_ims.get_affine())

        nb.save(all_ims, self.dwi_base+'.nii.gz')

        indices = ['1' for i in xrange(len(mux_ims1))] + [str(len(nondwi1)+1) for i in xrange(len(mux_ims2))]
        with open(self.index_file, 'w') as f:
            f.write(' '.join(indices))

        bvals = np.concatenate((bvals1[self.num_cal1:],bvals2[self.num_cal2:]), axis=0)
        bvecs = np.concatenate((bvecs1[:,self.num_cal1:],bvecs2[:,self.num_cal2:]), axis=1)
        with open(self.bval_file, 'w') as f:
            f.write(' '.join(['%0.1f' % value for value in bvals]))
        with open(self.bvec_file, 'w') as f:
            f.write(' '.join(['%0.4f' % value for value in bvecs[0,:]]) + '\n')
            f.write(' '.join(['%0.4f' % value for value in bvecs[1,:]]) + '\n')
            f.write(' '.join(['%0.4f' % value for value in bvecs[2,:]]) + '\n')

def b0_average(in_dwi, in_bval, max_b=10.0, out_file=None):
    """
    A function that averages the *b0* volumes from a DWI dataset.
    As current dMRI data are being acquired with all b-values > 0.0,
    the *lowb* volumes are selected by specifying the parameter max_b.

    .. warning:: *b0* should be already registered (head motion artifact should
      be corrected).

    """
    import numpy as np
    import nibabel as nb
    import os.path as op

    if out_file is None:
        fname, ext = op.splitext(op.basename(in_dwi))
        if ext == ".gz":
            fname, ext2 = op.splitext(fname)
            ext = ext2 + ext
        out_file = op.abspath("%s_avg_b0%s" % (fname, ext))

    imgs = np.array(nb.four_to_three(nb.load(in_dwi)))
    bval = np.loadtxt(in_bval)
    index = np.argwhere(bval <= max_b).flatten().tolist()

    b0s = [im.get_data().astype(np.float32)
           for im in imgs[index]]
    b0 = np.average(np.array(b0s), axis=0)

    hdr = imgs[0].get_header().copy()
    hdr.set_data_shape(b0.shape)
    hdr.set_xyzt_units('mm')
    hdr.set_data_dtype(np.float32)
    nb.Nifti1Image(b0, imgs[0].get_affine(), hdr).to_filename(out_file)
    return out_file

def time_avg(in_file, index=[0], out_file=None):
    """
    Average the input time-series, selecting the indices given in index

    .. warning:: time steps should be already registered (corrected for
      head motion artifacts).

    """
    import numpy as np
    import nibabel as nb
    import os.path as op

    if out_file is None:
        fname, ext = op.splitext(op.basename(in_file))
        if ext == ".gz":
            fname, ext2 = op.splitext(fname)
            ext = ext2 + ext
        out_file = op.abspath("%s_baseline%s" % (fname, ext))

    index = np.atleast_1d(index).tolist()

    imgs = np.array(nb.four_to_three(nb.load(in_file)))[index]
    if len(index) == 1:
        data = imgs[0].get_data().astype(np.float32)
    else:
        data = np.average(np.array([im.get_data().astype(np.float32)
                                    for im in imgs]), axis=0)

    hdr = imgs[0].get_header().copy()
    hdr.set_data_shape(data.shape)
    hdr.set_xyzt_units('mm')
    hdr.set_data_dtype(np.float32)
    nb.Nifti1Image(data, imgs[0].get_affine(), hdr).to_filename(out_file)
    return out_file

def b0_average(in_dwi, in_bval, out_file=None):
    """
    A function that averages the *b0* volumes from a DWI dataset.

    .. warning:: *b0* should be already registered (head motion artifact should
      be corrected).

    """
    import numpy as np
    import nibabel as nb
    import os.path as op

    if out_file is None:
        fname, ext = op.splitext(op.basename(in_dwi))
        if ext == ".gz":
            fname, ext2 = op.splitext(fname)
            ext = ext2 + ext
        out_file = op.abspath("%s_avg_b0%s" % (fname, ext))

    imgs = np.array(nb.four_to_three(nb.load(in_dwi)))
    bval = np.loadtxt(in_bval)
    b0s = [im.get_data().astype(np.float32)
           for im in imgs[np.where(bval == 0)]]
    b0 = np.average(np.array(b0s), axis=0)

    hdr = imgs[0].get_header().copy()
    hdr.set_data_shape(b0.shape)
    hdr.set_xyzt_units('mm')
    hdr.set_data_dtype(np.float32)
    nb.Nifti1Image(b0, imgs[0].get_affine(), hdr).to_filename(out_file)
    return out_file

    def _run_interface(self, runtime):
        mask_imgs = [nb.load(fname) for fname in self.inputs.label_files]
        if len(mask_imgs) == 1:
            mask_imgs = nb.four_to_three(mask_imgs[0])

        masks = [mask_img.get_data().astype(np.bool) for mask_img in mask_imgs]

        n_masks = len(masks)

        if n_masks != len(self.inputs.class_labels):
            raise ValueError("Number of masks must match number of labels")

        img = nb.load(self.inputs.in_file)

        series = np.zeros((img.shape[3], n_masks))

        data = img.get_data()
        for j in range(n_masks):
            series[:, j] = data[masks[j], :].mean(axis=0)

        output = np.vstack((self.inputs.class_labels, series.astype(str)))
        self._results['out_file'] = os.path.join(runtime.cwd,
                                                 self.inputs.out_file)
        np.savetxt(
            self._results['out_file'], output, fmt=b'%s', delimiter='\t')

        return runtime

def test_load_4D_img():
    # setup
    output_dir = os.path.join(OUTPUT_DIR, inspect.stack()[0][3])
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # try loading from 4D niimg
    film = create_random_image(n_scans=10)
    loaded_4D_img = load_4D_img(film)
    assert_true(is_niimg(loaded_4D_img))
    assert_equal(loaded_4D_img.shape, film.shape)

    # try loading from 4D image file
    film = create_random_image(n_scans=10)
    saved_img_filename = os.path.join(output_dir, "4D.nii.gz")
    nibabel.save(film, saved_img_filename)
    loaded_4D_img = load_4D_img(saved_img_filename)
    assert_true(is_niimg(loaded_4D_img))
    assert_equal(loaded_4D_img.shape, film.shape)

    # try loading from list of 3D niimgs
    film = create_random_image(n_scans=10)
    loaded_4D_img = load_4D_img(nibabel.four_to_three(film))
    assert_true(is_niimg(loaded_4D_img))
    assert_equal(loaded_4D_img.shape, film.shape)

    # try loading from list of 3D image files
    film = create_random_image(n_scans=10)
    saved_vols_filenames = save_vols(film,
                                     output_dir,
                                     ext='.nii.gz',
                                     )
    loaded_4D_img = load_4D_img(saved_vols_filenames)
    assert_true(is_niimg(loaded_4D_img))
    assert_equal(loaded_4D_img.shape, film.shape)

def test_save_vols():
    # setup
    n_scans = 10
    output_dir = os.path.join(OUTPUT_DIR, inspect.stack()[0][3])
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # create 4D film
    film = create_random_image(ndim=4, n_scans=n_scans)
    threeD_vols = nibabel.four_to_three(film)

    # save vols manually
    film_filename = os.path.join(output_dir, "film.nii.gz")
    threeD_vols_filenames = [os.path.join(output_dir, "fMETHODS-%06i" % i) for i in range(len(threeD_vols))]

    # check saving seperate 3D vols
    for stuff in [film, threeD_vols]:
        if isinstance(stuff, list):
            basenames = [os.path.basename(x) for x in threeD_vols_filenames]
        else:
            basenames = os.path.basename(film_filename)
        for concat in [False, True]:
            for bn in [None, basenames]:
                saved_vols_filenames = save_vols(stuff, output_dir, ext=".nii.gz", concat=concat, basenames=bn)
                if not concat and isinstance(stuff, list):
                    assert_true(isinstance(saved_vols_filenames, list))
                    assert_equal(len(saved_vols_filenames), n_scans)
                    if not bn is None:
                        assert_equal(os.path.basename(saved_vols_filenames[7]), "fMETHODS-000007.nii.gz")
                else:
                    assert_true(isinstance(saved_vols_filenames, basestring))
                    assert_true(saved_vols_filenames.endswith(".nii.gz"), msg=saved_vols_filenames)
                    assert_true(is_4D(check_niimg_4d(saved_vols_filenames)))

def read_niftis(file_lists):
    """
    Read niftis.

    Parameters
    ----------
    file_lists: list of list of paths.
        Each list of file paths is a unique class.

    Returns
    -------
    data, labels: tuple of array-like and list
        The data and corresponding labels
    """

    data0 = load_image(file_lists[0][0]).get_data()
    if len(data0.shape) == 3:
        x, y, z = data0.shape
        t = 1
    elif len(data0.shape) == 4:
        x, y, z, t = data0.shape
    else:
        raise ValueError("Cannot parse data with dimensions %r" % data0.shape)

    dt = (sum(len(fl) for fl in file_lists)) * t
    data = np.zeros((dt, x, y, z))

    labels = [[i] * (len(fl) * t) for i, fl in enumerate(file_lists)]
    labels = [item for sublist in labels for item in sublist]

    for i, fl in enumerate(file_lists):
        assert len([j for j in labels if j == i]) == len(fl) * t

    flattened_list = [item for sublist in file_lists for item in sublist]
    for i, f in enumerate(flattened_list):
        logger.info("Loading subject from file: %s%s" % (f, '' * 30))

        nifti = load_image(f)
        subject_data = nifti.get_data()

        if len(subject_data.shape) == 3:
            data[i] = subject_data
        elif len(subject_data.shape) == 4:
            data[i * t: (i + 1) * t] = subject_data.transpose((3, 0, 1, 2))
        else:
            raise ValueError("Cannot parse subject data with dimensions %r"
                             % subject_data.shape)

    logger.info("\rLoading subject from file: %s\n" % ('DONE' + " "*30))
    if data.shape[0] != len(labels):
        raise ValueError("Data and labels have different number of samples.")

    base_file = flattened_list[0]
    # Use nibabel in case we need to convert from 4d to 3d
    base = nib.load(base_file)
    if len(base.shape) == 4:
        base = nib.four_to_three(base)[0]

    return data, labels, base

    def _run_interface(self, runtime):
        if not self._have_nipy:
            raise RuntimeError('nipy is not installed')

        from nipy.algorithms.registration.histogram_registration import HistogramRegistration
        from nipy.algorithms.registration.affine import Affine

        vol1_nii = nb.load(self.inputs.volume1)
        vol2_nii = nb.load(self.inputs.volume2)

        dims = vol1_nii.get_data().ndim

        if dims == 3 or dims == 2:
            vols1 = [vol1_nii]
            vols2 = [vol2_nii]
        if dims == 4:
            vols1 = nb.four_to_three(vol1_nii)
            vols2 = nb.four_to_three(vol2_nii)

        if dims < 2 or dims > 4:
            raise RuntimeError('Image dimensions not supported (detected %dD file)' % dims)

        if isdefined(self.inputs.mask1):
            mask1 = nb.load(self.inputs.mask1).get_data() == 1
        else:
            mask1 = None

        if isdefined(self.inputs.mask2):
            mask2 = nb.load(self.inputs.mask2).get_data() == 1
        else:
            mask2 = None

        self._similarity = []

        for ts1, ts2 in zip(vols1, vols2):
            histreg = HistogramRegistration(from_img=ts1,
                                            to_img=ts2,
                                            similarity=self.inputs.metric,
                                            from_mask=mask1,
                                            to_mask=mask2)
            self._similarity.append(histreg.eval(Affine()))

        return runtime

def test_load_specific_vol():
    # setup
    output_dir = os.path.join(OUTPUT_DIR, inspect.stack()[0][3])
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    n_scans = 23

    # create 4D film
    film = create_random_image(ndim=4, n_scans=n_scans)

    # test loading vol from nibabel image object
    for t in xrange(n_scans):
        _vol, _n_scans = load_specific_vol(film, t)
        assert_equal(_n_scans, n_scans)
        assert_true(isinstance(_vol, type(film)))
        assert_equal(_vol.shape, film.shape[:-1])
        numpy.testing.assert_array_equal(_vol.get_data(),
                                         film.get_data()[..., t])

    # test loading vol from a single 4D filename
    for ext in IMAGE_EXTENSIONS:
        for film_filename_type in ['str', 'list']:
            if film_filename_type == 'str':
                # save film as single filename with extension ext
                film_filename = os.path.join(output_dir, "4D%s" % ext)
                nibabel.save(film, film_filename)
            else:
                # save film as multiple filenames (3D vols), with ext extension
                vols = nibabel.four_to_three(film)
                film_filename = []
                for t, vol in zip(xrange(n_scans), vols):
                    vol_filename = os.path.join(output_dir,
                                                "vol_%i%s" % (t, ext))
                    nibabel.save(vol, vol_filename)
                    film_filename.append(vol_filename)

            # test loading proper
            for t in xrange(n_scans):
                # note that .img loads as Nifti1Pair, not Nifti1Image
                vol_type = nibabel.Nifti1Pair if ext == '.img' else \
                    nibabel.Nifti1Image

                # load specific 3D vol from 4D film by filename
                _vol, _n_scans = load_specific_vol(film_filename, t)
                assert_equal(_n_scans, n_scans)
                assert_true(isinstance(_vol, vol_type))
                assert_equal(_vol.shape, film.shape[:-1])
                numpy.testing.assert_array_equal(_vol.get_data(),
                                              film.get_data()[..., t])

def test_do_3Dto4D_merge():
    # setup
    n_scans = 10
    output_dir = os.path.join(OUTPUT_DIR, inspect.stack()[0][3])
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # create 4D film
    film = create_random_image(ndim=4, n_scans=n_scans)
    threeD_vols = nibabel.four_to_three(film)

    _film = do_3Dto4D_merge(threeD_vols)

    assert_equal(_film.shape, film.shape)

    save_vols(threeD_vols, output_dir, ext='.nii.gz')

    def _run_interface(self, runtime):
        from dipy.reconst.csdeconv import ConstrainedSphericalDeconvModel
        from dipy.data import get_sphere
        # import marshal as pickle
        import pickle as pickle
        import gzip

        img = nb.load(self.inputs.in_file)
        imref = nb.four_to_three(img)[0]
        affine = img.get_affine()

        if isdefined(self.inputs.in_mask):
            msk = nb.load(self.inputs.in_mask).get_data()
        else:
            msk = np.ones(imref.get_shape())

        data = img.get_data().astype(np.float32)
        hdr = imref.get_header().copy()

        gtab = self._get_gradient_table()
        resp_file = np.loadtxt(self.inputs.response)

        response = (np.array(resp_file[0:3]), resp_file[-1])
        ratio = response[0][1] / response[0][0]

        if abs(ratio - 0.2) > 0.1:
            IFLOGGER.warn(('Estimated response is not prolate enough. '
                           'Ratio=%0.3f.') % ratio)

        csd_model = ConstrainedSphericalDeconvModel(
            gtab, response, sh_order=self.inputs.sh_order)

        IFLOGGER.info('Fitting CSD model')
        csd_fit = csd_model.fit(data, msk)

        f = gzip.open(self._gen_filename('csdmodel', ext='.pklz'), 'wb')
        pickle.dump(csd_model, f, -1)
        f.close()

        if self.inputs.save_fods:
            sphere = get_sphere('symmetric724')
            fods = csd_fit.odf(sphere)
            nb.Nifti1Image(fods.astype(np.float32), img.get_affine(),
                           None).to_filename(self._gen_filename('fods'))

        return runtime

def smooth_image(img, fwhm, **kwargs):
    """Function wrapper LinearFilter class. Spatially smoothens img with
    kernel of size fwhm.

    Parameters
    ----------
    img: ``ni.Nifti1Image``
        image to be smoothen
    fwhm: 1D array like of size as big as there are spatial dimensions
    in the image
        FWHM of smoothing kernel
    **kwargs: dict-like
        key-word arguments passed to LinearFilter constructor.

    Returns
    -------
    Smoothened image, same type and size as the input img.

    """

    if isinstance(img, basestring):
        img = ni.load(img)
    elif isinstance(img, tuple):
        assert len(img) == 2
        return smooth_image(ni.Nifti1Image(img[0], img[1]), fwhm, **kwargs)
    elif isinstance(img, list):
        return [smooth_image(x, fwhm, **kwargs) for x in img]
    else:
        assert is_niimg(img)

    if len(img.shape) == 4:
        return ni.concat_images(
            [smooth_image(vol, fwhm, **kwargs)
             for vol in ni.four_to_three(img)])
    else:
        assert len(img.shape) == 3

        smoothing_kernel = LinearFilter(
            img.get_affine(),
            img.shape,
            fwhm=fwhm,
            **kwargs)

        return ni.Nifti1Image(smoothing_kernel.smooth(img.get_data(),
                                                      clean=True),
                              img.get_affine())

def split_warp_volumes_fn(in_file):
    from nipype import logging
    from nipype.utils.filemanip import split_filename
    import nibabel as nb
    import os.path as op
    iflogger = logging.getLogger('interface')
    iflogger.info(in_file)
    path, name, ext = split_filename(in_file)
    image = nb.load(in_file)
    x_img, y_img, z_img = nb.four_to_three(image)
    x = op.abspath(name + '_x' + ".nii.gz")
    y = op.abspath(name + '_y' + ".nii.gz")
    z = op.abspath(name + '_z' + ".nii.gz")
    nb.save(x_img, x)
    nb.save(y_img, y)
    nb.save(z_img, z)
    return x, y, z

def test_transform():
    # setup
    output_dir = os.path.join(OUTPUT_DIR, inspect.stack()[0][3])
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    film = nibabel.Nifti1Image(np.random.rand(11, 13, 17, 19),
                               np.eye(4))
    threeD_vols = nibabel.four_to_three(film)

    # filenames
    film_filename = os.path.join(output_dir, 'film.nii.gz')
    threeD_vols_filenames = [os.path.join(output_dir, 'fMETHODS-%06i' % i)
                             for i in xrange(len(threeD_vols))]

    for stuff in [film, threeD_vols]:
        for as_files in [False, True]:
            if as_files:
                if isinstance(stuff, list):
                    basenames = [os.path.basename(x)
                                 for x in threeD_vols_filenames]
                else:
                    basenames = os.path.basename(film_filename)
                stuff = save_vols(stuff, output_dir, basenames=basenames)

            fmristc = fMRISTC().fit(raw_data=stuff)

            output = fmristc.transform(output_dir=output_dir)

            # test output type, shape, etc.
            if isinstance(stuff, list):
                nose.tools.assert_true(isinstance(
                        output, list))
                nose.tools.assert_equal(len(output),
                                        film.shape[-1])

                if as_files:
                    nose.tools.assert_equal(os.path.basename(output[7]),
                                            'afMETHODS-000007.nii.gz')
            else:
                if as_files:
                    nose.tools.assert_equal(os.path.basename(output),
                                            'afilm.nii.gz')                    

def load_specific_vol(vols, t, strict=False):
    """
    Utility function for loading specific volume on demand.

    Parameters
    ----------
    vols: string(s) or nibabel image object(s)
        input volumes or single 4D film
    t: int
        index of requested volume in the film

    """

    assert t >= 0

    if isinstance(vols, list):
        n_scans = len(vols)
        vol = load_vol(vols[t])
    elif is_niimg(vols) or isinstance(vols, basestring):
        _vols = nibabel.load(vols) if isinstance(vols, basestring) else vols
        if len(_vols.shape) != 4:
            if strict:
                raise ValueError(
                    "Expecting 4D image, got %iD" % len(_vols.shape))
            else:
                return _vols, 1

        n_scans = _vols.shape[-1]
        vol = nibabel.four_to_three(_vols)[t]
    else:  # unhandled type
        raise TypeError(
            ("vols must be string, image object, or list of such; "
             "got %s" % type(vols)))

    # delete trivial dimension
    if len(vol.shape) == 4:
        vol = nibabel.Nifti1Image(vol.get_data()[..., ..., ..., 0],
                                  vol.get_affine())

    assert is_niimg(vol)
    assert is_3D(vol)

    return vol, n_scans

def _load_mixed_gambles(zmap_imgs):
    """Ravel zmaps (one per subject) along time axis, resulting,
    in a n_subjects * n_trials 3D niimgs and, and then make
    gain vector y of same length.
    """
    X = []
    y = []
    mask = []
    for zmap_img in zmap_imgs:
        # load subject data
        this_X = zmap_img.get_data()
        affine = zmap_img.get_affine()
        finite_mask = np.all(np.isfinite(this_X), axis=-1)
        this_mask = np.logical_and(np.all(this_X != 0, axis=-1),
                                   finite_mask)
        this_y = np.array([np.arange(1, 9)] * 6).ravel()

        # gain levels
        if len(this_y) != this_X.shape[-1]:
            raise RuntimeError("%s: Expecting %i volumes, got %i!" % (
                zmap_img, len(this_y), this_X.shape[-1]))

        # standardize subject data
        this_X -= this_X.mean(axis=-1)[..., np.newaxis]
        std = this_X.std(axis=-1)
        std[std == 0] = 1
        this_X /= std[..., np.newaxis]

        # commit subject data
        X.append(this_X)
        y.extend(this_y)
        mask.append(this_mask)
    y = np.array(y)
    X = np.concatenate(X, axis=-1)
    mask = np.sum(mask, axis=0) > .5 * len(zmap_imgs)
    mask = np.logical_and(mask, np.all(np.isfinite(X), axis=-1))
    X = X[mask, :].T
    tmp = np.zeros(list(mask.shape) + [len(X)])
    tmp[mask, :] = X.T
    mask_img = nibabel.Nifti1Image(mask.astype(np.int), affine)
    X = nibabel.four_to_three(nibabel.Nifti1Image(tmp, affine))
    return X, y, mask_img

def _flatten_split_merge(in_files):
    from builtins import bytes, str

    if isinstance(in_files, (bytes, str)):
        in_files = [in_files]

    nfiles = len(in_files)

    all_nii = []
    for fname in in_files:
        nii = nb.squeeze_image(nb.load(fname))

        if nii.get_data().ndim > 3:
            all_nii += nb.four_to_three(nii)
        else:
            all_nii.append(nii)

    if len(all_nii) == 1:
        LOGGER.warn('File %s cannot be split', all_nii[0])
        return in_files[0], in_files

    if len(all_nii) == nfiles:
        flat_split = in_files
    else:
        splitname = genfname(in_files[0], suffix='split%04d')
        flat_split = []
        for i, nii in enumerate(all_nii):
            flat_split.append(splitname % i)
            nii.to_filename(flat_split[-1])

    # Only one 4D file was supplied
    if nfiles == 1:
        merged = in_files[0]
    else:
        # More that one in_files - need merge
        merged = genfname(in_files[0], suffix='merged')
        nb.concat_images(all_nii).to_filename(merged)

    return merged, flat_split

def remove_bad_volumes(dwi, bvec_file, bval_file, thresh=0.8):
    import numpy as np
    import nibabel as nb
    import os.path as op
    from nipype.utils.filemanip import split_filename

    dwi_4D = nb.load(dwi)
    dwi_files = nb.four_to_three(dwi_4D)

    bvecs = np.transpose(np.loadtxt(bvec_file))
    bvals = np.transpose(np.loadtxt(bval_file))

    bad_indices = np.where(np.abs(bvecs[:,0]) >= thresh)[0]
    n_removed = len(bad_indices)

    dwi_files = [i for j, i in enumerate(dwi_files) if j not in bad_indices]
    bvecs = [i for j, i in enumerate(bvecs) if j not in bad_indices]
    bvals = [i for j, i in enumerate(bvals) if j not in bad_indices]

    corr_dwi = nb.concat_images(dwi_files)
    corr_bvecs = np.transpose(bvecs)
    corr_bvals = np.transpose(bvals)

    assert(len(dwi_files) == len(bvecs) == len(bvals))

    _, name, _ = split_filename(dwi)
    out_dwi = op.abspath(name + "_vib.nii.gz")

    _, bvec_name, _ = split_filename(bvec_file)
    out_bvecs = op.abspath(bvec_name + "_vib.bvec")

    _, bval_name, _ = split_filename(bval_file)
    out_bvals = op.abspath(bval_name + "_vib.bval")

    nb.save(corr_dwi, out_dwi)
    np.savetxt(out_bvecs, corr_bvecs)
    np.savetxt(out_bvals, corr_bvals)
    print("%d volumes were removed at threshold %f" % (n_removed, thresh))
    return out_dwi, out_bvecs, out_bvals, n_removed