"""Unit tests for image iteration


"""
import logging
import unittest

import numpy

from data_models.polarisation import PolarisationFrame

from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter

from processing_components.image.operations import create_empty_image_like
from processing_components.simulation.testing_support import create_test_image

log = logging.getLogger(__name__)


class TestImageIterators(unittest.TestCase):
    def test_raster(self):
        
        m31original = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
        assert numpy.max(numpy.abs(m31original.data)), "Original is empty"
        
        for nraster in [1, 2, 4, 8, 9]:
            m31model = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
            for patch in image_raster_iter(m31model, facets=nraster):
                assert patch.data.shape[3] == (m31model.data.shape[3] // nraster), \
                    "Number of pixels in each patch: %d not as expected: %d" % (patch.data.shape[3],
                                                                                (m31model.data.shape[3] // nraster))
                assert patch.data.shape[2] == (m31model.data.shape[2] // nraster), \
                    "Number of pixels in each patch: %d not as expected: %d" % (patch.data.shape[2],
                                                                                (m31model.data.shape[2] // nraster))
                patch.data *= 2.0
            
            diff = m31model.data - 2.0 * m31original.data
            assert numpy.max(numpy.abs(m31model.data)), "Raster is empty for %d" % nraster
            assert numpy.max(numpy.abs(diff)) == 0.0, "Raster set failed for %d" % nraster

    def test_raster_exception(self):
    
        m31original = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
        assert numpy.max(numpy.abs(m31original.data)), "Original is empty"
    
        for nraster, overlap in [(-1, -1), (-1, 0), (2, 128), (1e6, 127)]:
            
            with self.assertRaises(AssertionError) as context:
                m31model = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
                for patch in image_raster_iter(m31model, facets=nraster, overlap=overlap):
                    patch.data *= 2.0
            
    def test_raster_overlap(self):
    
        m31original = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
        assert numpy.max(numpy.abs(m31original.data)), "Original is empty"
        flat = create_empty_image_like(m31original)
    
        for nraster, overlap in [(1, 0),  (1, 16), (4, 8), (4, 16), (8, 8), (16, 4), (9, 5)]:
            m31model = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
            for patch, flat_patch in zip(image_raster_iter(m31model, facets=nraster, overlap=overlap),
                                         image_raster_iter(flat, facets=nraster, overlap=overlap)):
                patch.data *= 2.0
                flat_patch.data[...] += 1.0
        
            assert numpy.max(numpy.abs(m31model.data)), "Raster is empty for %d" % nraster

    def test_raster_overlap_linear(self):
    
        m31original = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
        assert numpy.max(numpy.abs(m31original.data)), "Original is empty"
        flat = create_empty_image_like(m31original)
    
        for nraster, overlap in [(1, 0), (1, 16), (4, 8), (4, 16), (8, 8), (16, 4), (9, 5)]:
            m31model = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
            for patch, flat_patch in zip(image_raster_iter(m31model, facets=nraster, overlap=overlap,
                                                           taper='linear'),
                                         image_raster_iter(flat, facets=nraster, overlap=overlap)):
                patch.data *= 2.0
                flat_patch.data[...] += 1.0
        
            assert numpy.max(numpy.abs(m31model.data)), "Raster is empty for %d" % nraster

    def test_raster_overlap_quadratic(self):
    
        m31original = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
        assert numpy.max(numpy.abs(m31original.data)), "Original is empty"
        flat = create_empty_image_like(m31original)
    
        for nraster, overlap in [(1, 0), (1, 16), (4, 8), (4, 16), (8, 8), (16, 4), (9, 5)]:
            m31model = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
            for patch, flat_patch in zip(image_raster_iter(m31model, facets=nraster, overlap=overlap,
                                                           taper='quadratic'),
                                         image_raster_iter(flat, facets=nraster, overlap=overlap)):
                patch.data *= 2.0
                flat_patch.data[...] += 1.0
        
            assert numpy.max(numpy.abs(m31model.data)), "Raster is empty for %d" % nraster

    def test_raster_overlap_tukey(self):
    
        m31original = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
        assert numpy.max(numpy.abs(m31original.data)), "Original is empty"
        flat = create_empty_image_like(m31original)
    
        for nraster, overlap in [(1, 0), (1, 16), (4, 8), (4, 16), (8, 8), (16, 4), (9, 5)]:
            m31model = create_test_image(polarisation_frame=PolarisationFrame('stokesI'))
            for patch, flat_patch in zip(image_raster_iter(m31model, facets=nraster, overlap=overlap,
                                                           taper='tukey'),
                                         image_raster_iter(flat, facets=nraster, overlap=overlap)):
                patch.data *= 2.0
                flat_patch.data[...] += 1.0
        
            assert numpy.max(numpy.abs(m31model.data)), "Raster is empty for %d" % nraster

    def test_channelise(self):
        m31cube = create_test_image(polarisation_frame=PolarisationFrame('stokesI'),
                                        frequency=numpy.linspace(1e8,1.1e8, 128))
        
        for subimages in [128, 16, 8, 2, 1]:
            for slab in image_channel_iter(m31cube, subimages=subimages):
                assert slab.data.shape[0] == 128 // subimages

    def test_null(self):
        m31cube = create_test_image(polarisation_frame=PolarisationFrame('stokesI'),
                                    frequency=numpy.linspace(1e8, 1.1e8, 128))
    
        for i, im in enumerate(image_null_iter(m31cube)):
            assert i<1, "Null iterator returns more than one value"


if __name__ == '__main__':
    unittest.main()

[(17, 'logging.getLogger', 'logging.getLogger', False, 'import logging\n'), (133, 'unittest.main', 'unittest.main', False, 'import unittest\n'), (57, 'processing_components.image.operations.create_empty_image_like', 'create_empty_image_like', False, 'from processing_components.image.operations import create_empty_image_like\n'), (72, 'processing_components.image.operations.create_empty_image_like', 'create_empty_image_like', False, 'from processing_components.image.operations import create_empty_image_like\n'), (88, 'processing_components.image.operations.create_empty_image_like', 'create_empty_image_like', False, 'from processing_components.image.operations import create_empty_image_like\n'), (104, 'processing_components.image.operations.create_empty_image_like', 'create_empty_image_like', False, 'from processing_components.image.operations import create_empty_image_like\n'), (24, 'numpy.abs', 'numpy.abs', False, 'import numpy\n'), (28, 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), (44, 'numpy.abs', 'numpy.abs', False, 'import numpy\n'), (56, 'numpy.abs', 'numpy.abs', False, 'import numpy\n'), (71, 'numpy.abs', 'numpy.abs', False, 'import numpy\n'), (87, 'numpy.abs', 'numpy.abs', False, 'import numpy\n'), (103, 'numpy.abs', 'numpy.abs', False, 'import numpy\n'), (121, 'processing_components.image.iterators.image_channel_iter', 'image_channel_iter', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), (128, 'processing_components.image.iterators.image_null_iter', 'image_null_iter', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), (23, 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', False, 'from data_models.polarisation import PolarisationFrame\n'), (38, 'numpy.abs', 'numpy.abs', False, 'import numpy\n'), (43, 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', False, 'from data_models.polarisation import PolarisationFrame\n'), (50, 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), (55, 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', False, 'from data_models.polarisation import PolarisationFrame\n'), (61, 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), (62, 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), (66, 'numpy.abs', 'numpy.abs', False, 'import numpy\n'), (70, 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', False, 'from data_models.polarisation import PolarisationFrame\n'), (76, 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), (78, 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), (82, 'numpy.abs', 'numpy.abs', False, 'import numpy\n'), (86, 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', False, 'from data_models.polarisation import PolarisationFrame\n'), (92, 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), (94, 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), (98, 'numpy.abs', 'numpy.abs', False, 'import numpy\n'), (102, 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', False, 'from data_models.polarisation import PolarisationFrame\n'), (108, 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), (110, 'processing_components.image.iterators.image_raster_iter', 'image_raster_iter', False, 'from processing_components.image.iterators import image_raster_iter, image_channel_iter, image_null_iter\n'), (114, 'numpy.abs', 'numpy.abs', False, 'import numpy\n'), (117, 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', False, 'from data_models.polarisation import PolarisationFrame\n'), (118, 'numpy.linspace', 'numpy.linspace', False, 'import numpy\n'), (125, 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', False, 'from data_models.polarisation import PolarisationFrame\n'), (126, 'numpy.linspace', 'numpy.linspace', False, 'import numpy\n'), (27, 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', False, 'from data_models.polarisation import PolarisationFrame\n'), (39, 'numpy.abs', 'numpy.abs', False, 'import numpy\n'), (60, 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', False, 'from data_models.polarisation import PolarisationFrame\n'), (75, 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', False, 'from data_models.polarisation import PolarisationFrame\n'), (91, 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', False, 'from data_models.polarisation import PolarisationFrame\n'), (107, 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', False, 'from data_models.polarisation import PolarisationFrame\n'), (49, 'data_models.polarisation.PolarisationFrame', 'PolarisationFrame', False, 'from data_models.polarisation import PolarisationFrame\n')]