"""
Test that any previously reported bugs are fixed

nptest.assert_allclose
self.assertEqual
self.assertTrue

"""

import os
import unittest
from common import run_tests

import ants

import numpy as np
import numpy.testing as nptest


class Test_bugs(unittest.TestCase):
    """
    Test ants.ANTsImage class
    """
    def setUp(self):
        pass

    def tearDown(self):
        pass

    def test_resample_returns_NaNs(self):
        """
        Test that resampling an image doesnt cause the resampled
        image to have NaNs - previously caused by resampling an
        image of type DOUBLE
        """
        img2d = ants.image_read(ants.get_ants_data('r16'))
        img2dr = ants.resample_image(img2d, (2,2), 0, 0)

        self.assertTrue(np.sum(np.isnan(img2dr.numpy())) == 0)

        img3d = ants.image_read(ants.get_ants_data('mni'))
        img3dr = ants.resample_image(img3d, (2,2,2), 0, 0)

        self.assertTrue(np.sum(np.isnan(img3dr.numpy())) == 0)

    def test_compose_multi_type_transforms(self):
        image = ants.image_read(ants.get_ants_data("r16"))

        linear_transform = ants.create_ants_transform(transform_type=
            "AffineTransform", precision='float', dimension=image.dimension)

        displacement_field = ants.simulate_displacement_field(image,
            field_type="bspline", number_of_random_points=1000,
            sd_noise=10.0, enforce_stationary_boundary=True,
            number_of_fitting_levels=4, mesh_size=1,
            sd_smoothing=4.0)
        displacement_field_xfrm = ants.transform_from_displacement_field(displacement_field)

        xfrm = ants.compose_ants_transforms([linear_transform, displacement_field_xfrm])
        xfrm = ants.compose_ants_transforms([linear_transform, linear_transform])
        xfrm = ants.compose_ants_transforms([displacement_field_xfrm, linear_transform])

    def test_bspline_image_with_2d_weights(self):
        # see https://github.com/ANTsX/ANTsPy/issues/655
        import ants
        import numpy as np

        output_size = (256, 256)
        bspline_epsilon = 1e-4
        number_of_fitting_levels = 4

        image = ants.image_read(ants.get_ants_data("r16"))
        image = ants.resample_image(image, (100, 100), use_voxels=True)

        indices = np.meshgrid(list(range(image.shape[0])), 
                                list(range(image.shape[1])))
        indices_array = np.stack((indices[1].flatten(), 
                                indices[0].flatten()), axis=0)

        image_parametric_values = indices_array.transpose()

        weight_array = np.ones(image.shape)
        parametric_values = image_parametric_values
        scattered_data = np.atleast_2d(image.numpy().flatten()).transpose()
        weight_values = np.atleast_2d(weight_array.flatten()).transpose()

        min_parametric_values = np.min(parametric_values, axis=0)
        max_parametric_values = np.max(parametric_values, axis=0)

        spacing = np.zeros((2,))
        for d in range(2):
            spacing[d] = (max_parametric_values[d] - min_parametric_values[d]) / (output_size[d] - 1) + bspline_epsilon

        bspline_image = ants.fit_bspline_object_to_scattered_data(scattered_data, parametric_values, 
                                                                parametric_domain_origin=min_parametric_values - bspline_epsilon,
                                                                parametric_domain_spacing=spacing,
                                                                parametric_domain_size=output_size,
                                                                data_weights=weight_values,
                                                                number_of_fitting_levels=number_of_fitting_levels,
                                                                mesh_size=1)


    def test_scalar_rgb_missing(self):
        import ants
        img = ants.image_read(ants.get_data('r16'))
        with self.assertRaises(Exception):
            img_color = ants.scalar_to_rgb(img, cmap='jet')

    def test_bspline_zeros(self):
        import ants
        import numpy as np
        x = np.linspace(-4, 4, num=100) + np.random.uniform(-0.1, 0.1, 100)
        u = np.linspace(0, 1.0, num=len(x))
        scattered_data = np.expand_dims(u, axis=-1)
        parametric_data = np.expand_dims(u, axis=-1)
        spacing = 1/(len(x)-1) * 1.0
        bspline_curve = ants.fit_bspline_object_to_scattered_data(
                scattered_data, parametric_data,
                parametric_domain_origin=[0.0], parametric_domain_spacing=[spacing],
                parametric_domain_size=[len(x)], is_parametric_dimension_closed=None,
                number_of_fitting_levels=5, mesh_size=1)

        # Erroneously returns all zeros.
        self.assertNotEqual(bspline_curve.sum(), 0)

    def test_from_numpy_different_dtypes(self):
        all_dtypes = ('bool',
                      'int8',
                      'int16',
                      'int32',
                      'int64',
                      'uint16',
                      'uint64',
                      'float16')
        arr = np.random.randn(100,100)
        
        for dtype in all_dtypes:
            arr2 = arr.astype(dtype)
            img = ants.from_numpy(arr2)
            self.assertTrue(ants.is_image(img))
        
if __name__ == '__main__':
    run_tests()