def test_resampling_continuous_with_affine():
    prng = np.random.RandomState(10)

    data_3d = prng.randint(1, 4, size=(1, 10, 10))
    data_4d = prng.randint(1, 4, size=(1, 10, 10, 3))

    for data in [data_3d, data_4d]:
        for angle in (0, np.pi / 2., np.pi, 3 * np.pi / 2.):
            rot = rotation(0, angle)

            img = Nifti1Image(data, np.eye(4))
            rot_img = resample_img(
                img,
                target_affine=rot,
                interpolation='continuous')
            rot_img_back = resample_img(
                rot_img,
                target_affine=np.eye(4),
                interpolation='continuous')

            center = slice(1, 9)
            # values on the edges are wrong for some reason
            mask = (0, center, center)
            np.testing.assert_allclose(
                img.get_data()[mask],
                rot_img_back.get_data()[mask])
            assert_equal(rot_img.get_data().dtype,
                         np.dtype(data.dtype.name.replace('int', 'float')))

def test_resampling_fill_value():
    """ Test resampling with a non-zero fill value
    """
    prng = np.random.RandomState(10)

    data_3d = prng.rand(1, 4, 4)
    data_4d = prng.rand(1, 4, 4, 3)

    angle = np.pi/4
    rot = rotation(0, angle)

    # Try a few different fill values
    for data in [data_3d, data_4d]:
        for val in (-3.75, 0):
            if val:
                rot_img = resample_img(Nifti1Image(data, np.eye(4)),
                                       target_affine=rot,
                                       interpolation='nearest',
                                       fill_value=val,
                                       clip=False)
            else:
                rot_img = resample_img(Nifti1Image(data, np.eye(4)),
                                       target_affine=rot,
                                       interpolation='nearest',
                                       clip=False)
            assert_equal(rot_img.get_data().flatten()[0],
                         val)

            rot_img2 = resample_to_img(Nifti1Image(data, np.eye(4)),
                                       rot_img,
                                       interpolation='nearest',
                                       fill_value=val)
            assert_equal(rot_img2.get_data().flatten()[0],
                         val)

def test_identity_resample():
    """ Test resampling with an identity affine.
    """
    shape = (3, 2, 5, 2)
    data = np.random.randint(0, 10, shape)
    affine = np.eye(4)
    affine[:3, -1] = 0.5 * np.array(shape[:3])
    rot_img = resample_img(Nifti1Image(data, affine),
                           target_affine=affine, interpolation='nearest')
    np.testing.assert_almost_equal(data, rot_img.get_data())
    # Smoke-test with a list affine
    rot_img = resample_img(Nifti1Image(data, affine),
                           target_affine=affine.tolist(),
                           interpolation='nearest')
    # Test with a 3x3 affine
    rot_img = resample_img(Nifti1Image(data, affine),
                           target_affine=affine[:3, :3],
                           interpolation='nearest')
    np.testing.assert_almost_equal(data, rot_img.get_data())

    # Test with non native endian data

    # Test with big endian data ('>f8')
    for interpolation in ['nearest', 'continuous']:
        rot_img = resample_img(Nifti1Image(data.astype('>f8'), affine),
                               target_affine=affine.tolist(),
                               interpolation=interpolation)
        np.testing.assert_almost_equal(data, rot_img.get_data())

    # Test with little endian data ('<f8')
    for interpolation in ['nearest', 'continuous']:
        rot_img = resample_img(Nifti1Image(data.astype('<f8'), affine),
                               target_affine=affine.tolist(),
                               interpolation=interpolation)
        np.testing.assert_almost_equal(data, rot_img.get_data())

def test_resampling_with_affine():
    """ Test resampling with a given rotation part of the affine.
    """
    prng = np.random.RandomState(10)

    data_3d = prng.randint(4, size=(1, 4, 4))
    data_4d = prng.randint(4, size=(1, 4, 4, 3))

    for data in [data_3d, data_4d]:
        for angle in (0, np.pi, np.pi / 2., np.pi / 4., np.pi / 3.):
            rot = rotation(0, angle)
            rot_img = resample_img(Nifti1Image(data, np.eye(4)),
                                   target_affine=rot,
                                   interpolation='nearest')
            assert_equal(np.max(data),
                         np.max(rot_img.get_data()))
            assert_equal(rot_img.get_data().dtype, data.dtype)

    # We take the same rotation logic as above and test with nonnative endian
    # data as input
    for data in [data_3d, data_4d]:
        img = Nifti1Image(data.astype('>f8'), np.eye(4))
        for angle in (0, np.pi, np.pi / 2., np.pi / 4., np.pi / 3.):
            rot = rotation(0, angle)
            rot_img = resample_img(img, target_affine=rot,
                                   interpolation='nearest')
            assert_equal(np.max(data),
                         np.max(rot_img.get_data()))

def save_maps(model_dir, doc, resample=False,
              target_affine=None, target_shape=None):
    for dtype in ['c_maps', 't_maps']:
        if dtype in doc:
            maps_dir = make_dir(model_dir, dtype, strict=False)
            for key in doc[dtype]:
                fname = '%s.nii.gz' % safe_name(key.lower())
                img = nb.load(doc[dtype][key])
                if resample:
                    img = resample_img(img, target_affine, target_shape)
                nb.save(img, os.path.join(maps_dir, fname))
    if 'beta_maps' in doc:
        maps_dir = make_dir(model_dir, 'beta_maps')
        for path in doc['beta_maps']:
            fname = '%s.nii.gz' % safe_name(os.path.split(
                path)[1].lower().split('.')[0])
            img = nb.load(path)
            if resample:
                img = resample_img(
                    img, target_affine, target_shape, copy=False)
            nb.save(img, os.path.join(maps_dir, fname))
    if 'mask' in doc:
        img = nb.load(doc['mask'])
        if resample:
            img = resample_img(img, target_affine, target_shape,
                               interpolation='nearest', copy=False)
        nb.save(img, os.path.join(model_dir, 'mask.nii.gz'))

def test_resampling_with_int_types_no_crash():
    affine = np.eye(4)
    data = np.zeros((2, 2, 2))

    for dtype in [np.int, np.int8, np.int16, np.int32, np.int64,
                  np.uint, np.uint8, np.uint16, np.uint32, np.uint64,
                  np.float32, np.float64, np.float]:
        img = Nifti1Image(data.astype(dtype), affine)
        resample_img(img, target_affine=2. * affine)

def test_resampling_error_checks():
    shape = (3, 2, 5, 2)
    target_shape = (5, 3, 2)
    affine = np.eye(4)
    data = np.random.randint(0, 10, shape)
    img = Nifti1Image(data, affine)

    # Correct parameters: no exception
    resample_img(img, target_shape=target_shape, target_affine=affine)
    resample_img(img, target_affine=affine)

    with testing.write_tmp_imgs(img) as filename:
        resample_img(filename, target_shape=target_shape, target_affine=affine)

    # Missing parameter
    assert_raises(ValueError, resample_img, img, target_shape=target_shape)

    # Invalid shape
    assert_raises(ValueError, resample_img, img, target_shape=(2, 3),
                  target_affine=affine)

    # Invalid interpolation
    interpolation = 'an_invalid_interpolation'
    pattern = "interpolation must be either.+{0}".format(interpolation)
    testing.assert_raises_regex(ValueError, pattern,
                                resample_img, img, target_shape=target_shape,
                                target_affine=affine,
                                interpolation="an_invalid_interpolation")

    # Noop
    target_shape = shape[:3]

    img_r = resample_img(img, copy=False)
    assert_equal(img_r, img)

    img_r = resample_img(img, copy=True)
    assert_false(np.may_share_memory(img_r.get_data(), img.get_data()))

    np.testing.assert_almost_equal(img_r.get_data(), img.get_data())
    np.testing.assert_almost_equal(img_r.get_affine(), img.get_affine())

    img_r = resample_img(img, target_affine=affine, target_shape=target_shape,
                         copy=False)
    assert_equal(img_r, img)

    img_r = resample_img(img, target_affine=affine, target_shape=target_shape,
                         copy=True)
    assert_false(np.may_share_memory(img_r.get_data(), img.get_data()))
    np.testing.assert_almost_equal(img_r.get_data(), img.get_data())
    np.testing.assert_almost_equal(img_r.get_affine(), img.get_affine())

def test_4d_affine_bounding_box_error():

    small_data = np.ones([4, 4, 4])
    small_data_4D_affine = np.eye(4)
    small_data_4D_affine[:3, -1] = np.array([5, 4, 5])

    small_img = Nifti1Image(small_data,
                            small_data_4D_affine)

    bigger_data_4D_affine = np.eye(4)
    bigger_data = np.zeros([10, 10, 10])
    bigger_img = Nifti1Image(bigger_data,
                             bigger_data_4D_affine)

    # We would like to check whether all/most of the data
    # will be contained in the resampled image
    # The measure will be the l2 norm, since some resampling
    # schemes approximately conserve it

    def l2_norm(arr):
        return (arr ** 2).sum()

    # resample using 4D affine and specified target shape
    small_to_big_with_shape = resample_img(
        small_img,
        target_affine=bigger_img.get_affine(),
        target_shape=bigger_img.shape)
    # resample using 3D affine and no target shape
    small_to_big_without_shape_3D_affine = resample_img(
        small_img,
        target_affine=bigger_img.get_affine()[:3, :3])
    # resample using 4D affine and no target shape
    small_to_big_without_shape = resample_img(
        small_img,
        target_affine=bigger_img.get_affine())

    # The first 2 should pass
    assert_almost_equal(l2_norm(small_data),
                 l2_norm(small_to_big_with_shape.get_data()))
    assert_almost_equal(l2_norm(small_data),
                 l2_norm(small_to_big_without_shape_3D_affine.get_data()))

    # After correcting decision tree for 4x4 affine given + no target shape
    # from "use initial shape" to "calculate minimal bounding box respecting
    # the affine anchor and the data"
    assert_almost_equal(l2_norm(small_data),
                 l2_norm(small_to_big_without_shape.get_data()))

    assert_array_equal(small_to_big_without_shape.shape,
                 small_data_4D_affine[:3, -1] + np.array(small_img.shape))

def test_resample_img_segmentation_fault():
    # see https://github.com/nilearn/nilearn/issues/346
    shape_in = (64, 64, 64)
    aff_in = np.diag([2., 2., 2., 1.])
    aff_out = np.diag([3., 3., 3., 1.])
    # fourth_dim = 1024 works fine but for 1025 creates a segmentation
    # fault with scipy < 0.14.1
    fourth_dim = 1025
    data = np.ones(shape_in + (fourth_dim, ), dtype=np.float64)

    img_in = Nifti1Image(data, aff_in)

    resample_img(img_in,
                 target_affine=aff_out,
                 interpolation='nearest')

def test_resampling_result_axis_permutation():
    # Transform real data using easily checkable transformations
    # For now: axis permutations
    # create a cuboid full of deterministic data, padded with one
    # voxel thickness of zeros
    core_shape = (3, 5, 4)
    core_data = np.arange(np.prod(core_shape)).reshape(core_shape)
    full_data_shape = np.array(core_shape) + 2
    full_data = np.zeros(full_data_shape)
    full_data[[slice(1, 1 + s) for s in core_shape]] = core_data

    source_img = Nifti1Image(full_data, np.eye(4))

    axis_permutations = [[0, 1, 2],
                         [1, 0, 2],
                         [2, 1, 0],
                         [0, 2, 1]]

    # check 3x3 transformation matrix
    for ap in axis_permutations:
        target_affine = np.eye(3)[ap]
        resampled_img = resample_img(source_img,
                                     target_affine=target_affine)

        resampled_data = resampled_img.get_data()
        what_resampled_data_should_be = full_data.transpose(ap)
        assert_array_almost_equal(resampled_data,
                                  what_resampled_data_should_be)

    # check 4x4 transformation matrix
    offset = np.array([-2, 1, -3])
    for ap in axis_permutations:
        target_affine = np.eye(4)
        target_affine[:3, :3] = np.eye(3)[ap]
        target_affine[:3, 3] = offset

        resampled_img = resample_img(source_img,
                                     target_affine=target_affine)
        resampled_data = resampled_img.get_data()
        offset_cropping = np.vstack([-offset[ap][np.newaxis, :],
                                     np.zeros([1, 3])]
                                    ).T.ravel().astype(int)
        what_resampled_data_should_be = pad(full_data.transpose(ap),
                                            *list(offset_cropping))

        assert_array_almost_equal(resampled_data,
                                  what_resampled_data_should_be)

def test_resample_img_segmentation_fault():
    if 'APPVEYOR_TEST' in os.environ:
        raise(SkipTest('Skipped on appveyor due to not enough memory on image'))

    # see https://github.com/nilearn/nilearn/issues/346
    shape_in = (64, 64, 64)
    aff_in = np.diag([2., 2., 2., 1.])
    aff_out = np.diag([3., 3., 3., 1.])
    # fourth_dim = 1024 works fine but for 1025 creates a segmentation
    # fault with scipy < 0.14.1
    fourth_dim = 1025
    data = np.ones(shape_in + (fourth_dim, ), dtype=np.float64)

    img_in = Nifti1Image(data, aff_in)

    resample_img(img_in,
                 target_affine=aff_out,
                 interpolation='nearest')

def test_identity_resample():
    """ Test resampling with an identity affine.
    """
    shape = (3, 2, 5, 2)
    data = np.random.randint(0, 10, shape)
    affine = np.eye(4)
    affine[:3, -1] = 0.5 * np.array(shape[:3])
    rot_img = resample_img(Nifti1Image(data, affine),
                           target_affine=affine, interpolation='nearest')
    np.testing.assert_almost_equal(data, rot_img.get_data())
    # Smoke-test with a list affine
    rot_img = resample_img(Nifti1Image(data, affine),
                           target_affine=affine.tolist(),
                           interpolation='nearest')
    # Test with a 3x3 affine
    rot_img = resample_img(Nifti1Image(data, affine),
                           target_affine=affine[:3, :3],
                           interpolation='nearest')
    np.testing.assert_almost_equal(data, rot_img.get_data())

def test_resample_img_segmentation_fault():
    # see https://github.com/nilearn/nilearn/issues/346
    shape_in = (64, 64, 64)
    aff_in = np.diag([2., 2., 2., 1.])
    aff_out = np.diag([3., 3., 3., 1.])
    # fourth_dim = 1024 works fine but for 1025 creates a segmentation
    # fault with scipy < 0.14.1
    fourth_dim = 1025

    try:
        data = np.ones(shape_in + (fourth_dim, ), dtype=np.float64)
    except MemoryError:
        # This can happen on AppVeyor and for 32-bit Python on Windows
        raise SkipTest('Not enough RAM to run this test')

    img_in = Nifti1Image(data, aff_in)

    resample_img(img_in,
                 target_affine=aff_out,
                 interpolation='nearest')

def linear_downsampling_generator(generator, max_downsampling_factor=2, isotropic=False):
    '''
    Downsamples each sample (linearly) by a random factor and upsamples to original resolution again (nearest neighbor).

    Info:
    * Uses nilearn resample_img for resampling.
    * If isotropic=True:  Resamples all dimensions (channels, x, y, z) with same downsampling factor
    * If isotropic=False: Randomly choose new downsampling factor for each dimension
    * Does not resample "seg".
    '''
    import nibabel as nib
    from nilearn.image.resampling import resample_img, resample_to_img

    for data_dict in generator:
        assert "data" in list(
            data_dict.keys()), "your data generator needs to return a python dictionary with at least a 'data' key value pair"

        data = data_dict[
            'data']  # shape of data must be: (batch_size, nr_of_channels, x, y, [z])  (z ist optional; nr_of_channels can be 1)
        dim = len(data.shape[2:])  # remove batch_size and nr_of_channels dimension
        for sample_idx in range(data.shape[0]):

            fact = random.uniform(1, max_downsampling_factor)

            for channel_idx in range(data.shape[1]):

                affine = np.identity(4)
                if dim == 3:
                    img_data = data[sample_idx, channel_idx]
                elif dim == 2:
                    tmp = data[sample_idx, channel_idx]
                    img_data = np.reshape(tmp, (
                    1, tmp.shape[0], tmp.shape[1]))  # add third spatial dimension to make resample_img work
                else:
                    raise ValueError("Invalid dimension size")

                image = nib.Nifti1Image(img_data, affine)
                affine2 = affine
                if isotropic:
                    affine2[0, 0] = fact
                    affine2[1, 1] = fact
                    affine2[2, 2] = fact
                else:
                    affine2[0, 0] = random.uniform(1, max_downsampling_factor)
                    affine2[1, 1] = random.uniform(1, max_downsampling_factor)
                    affine2[2, 2] = random.uniform(1, max_downsampling_factor)
                affine2[3, 3] = 1
                image2 = resample_img(image, target_affine=affine2, interpolation='continuous')
                image3 = resample_to_img(image2, image, 'nearest')
                data[sample_idx, channel_idx] = np.squeeze(image3.get_data())

        data_dict["data"] = data
        yield data_dict

def test_resampling_with_affine():
    """ Test resampling with a given rotation part of the affine.
    """
    prng = np.random.RandomState(10)
    data = prng.randint(4, size=(1, 4, 4))
    for angle in (0, np.pi, np.pi / 2., np.pi / 4., np.pi / 3.):
        rot = rotation(0, angle)
        rot_img = resample_img(Nifti1Image(data, np.eye(4)),
                               target_affine=rot,
                               interpolation='nearest')
        np.testing.assert_almost_equal(np.max(data),
                                       np.max(rot_img.get_data()))

def test_downsample():
    """ Test resampling with a 1/2 down-sampling affine.
    """
    rand_gen = np.random.RandomState(0)
    shape = (6, 3, 6, 2)
    data = rand_gen.random_sample(shape)
    affine = np.eye(4)
    rot_img = resample_img(Nifti1Image(data, affine),
                           target_affine=2 * affine, interpolation='nearest')
    downsampled = data[::2, ::2, ::2, ...]
    x, y, z = downsampled.shape[:3]
    np.testing.assert_almost_equal(downsampled,
                                   rot_img.get_data()[:x, :y, :z, ...])

def test_resample_clip():
    # Resample and image and get larger and smaller
    # value than in the original. Use clip to get rid of these images

    shape = (6, 3, 6)
    data = np.zeros(shape=shape)
    data[1:-2, 1:-1, 1:-2] = 1

    source_affine = np.diag((2, 2, 2, 1))
    source_img = Nifti1Image(data, source_affine)

    target_affine = np.eye(4)
    no_clip_data = resample_img(source_img, target_affine,
                                clip=False).get_data()
    clip_data = resample_img(source_img,
                             target_affine, clip=True).get_data()

    not_clip = np.where((no_clip_data > data.min()) & (no_clip_data < data.max()))

    assert_true(np.any(no_clip_data > data.max()))
    assert_true(np.any(no_clip_data < data.min()))
    assert_true(np.all(clip_data <= data.max()))
    assert_true(np.all(clip_data >= data.min()))
    assert_array_equal(no_clip_data[not_clip], clip_data[not_clip])

def test_downsample():
    """ Test resampling with a 1/2 down-sampling affine.
    """
    rand_gen = np.random.RandomState(0)
    shape = (6, 3, 6, 2)
    data = rand_gen.random_sample(shape)
    affine = np.eye(4)
    rot_img = resample_img(Nifti1Image(data, affine),
                           target_affine=2 * affine, interpolation='nearest')
    downsampled = data[::2, ::2, ::2, ...]
    x, y, z = downsampled.shape[:3]
    np.testing.assert_almost_equal(downsampled,
                                   rot_img.get_data()[:x, :y, :z, ...])

    # Test with non native endian data

    # Test to check that if giving non native endian data as input should
    # work as normal and expected to return the same output as above tests.

    # Big endian data ('>f8')
    for copy in [True, False]:
        rot_img = resample_img(Nifti1Image(data.astype('>f8'), affine),
                               target_affine=2 * affine,
                               interpolation='nearest',
                               copy=copy)
        np.testing.assert_almost_equal(downsampled,
                                       rot_img.get_data()[:x, :y, :z, ...])

    # Little endian data
    for copy in [True, False]:
        rot_img = resample_img(Nifti1Image(data.astype('<f8'), affine),
                               target_affine=2 * affine,
                               interpolation='nearest',
                               copy=copy)
        np.testing.assert_almost_equal(downsampled,
                                       rot_img.get_data()[:x, :y, :z, ...])

    def _get_resampled_img(dtype):
        data = np.ones((10, 10, 10), dtype=dtype)
        data[3:7, 3:7, 3:7] = 2

        affine = np.eye(4)

        theta = math.pi / 6.
        c = math.cos(theta)
        s = math.sin(theta)

        affine = np.array([[1, 0, 0, 0],
                           [0, c, -s, 0],
                           [0, s, c, 0],
                           [0, 0, 0, 1]])

        img = Nifti1Image(data, affine)
        return resample_img(img, target_affine=np.eye(4))

def augment_linear_downsampling_nilearn(data, max_downsampling_factor=2, isotropic=False):
    '''
    Downsamples each sample (linearly) by a random factor and upsamples to original resolution again (nearest neighbor).

    Info:
    * Uses nilearn resample_img for resampling.
    * If isotropic=True:  Resamples all channels (channels, x, y, z) with same downsampling factor
    * If isotropic=False: Randomly choose new downsampling factor for each dimension
    '''
    import nibabel as nib
    from nilearn.image.resampling import resample_img, resample_to_img

    dim = len(data.shape[2:])  # remove batch_size and nr_of_channels dimension
    for sample_idx in range(data.shape[0]):

        fact = random.uniform(1, max_downsampling_factor)

        for channel_idx in range(data.shape[1]):

            affine = np.identity(4)
            if dim == 3:
                img_data = data[sample_idx, channel_idx]
            elif dim == 2:
                tmp = data[sample_idx, channel_idx]
                img_data = np.reshape(tmp, (
                1, tmp.shape[0], tmp.shape[1]))  # add third spatial dimension to make resample_img work
            else:
                raise ValueError("Invalid dimension size")

            image = nib.Nifti1Image(img_data, affine)
            affine2 = affine
            if isotropic:
                affine2[0, 0] = fact
                affine2[1, 1] = fact
                affine2[2, 2] = fact
            else:
                affine2[0, 0] = random.uniform(1, max_downsampling_factor)
                affine2[1, 1] = random.uniform(1, max_downsampling_factor)
                affine2[2, 2] = random.uniform(1, max_downsampling_factor)
            affine2[3, 3] = 1
            image2 = resample_img(image, target_affine=affine2, interpolation='continuous')
            image3 = resample_to_img(image2, image, 'nearest')

        data[sample_idx, channel_idx] = np.squeeze(image3.get_data())
    return data