import numpy as np

from autolens.profiles import light_profiles
from profiling import profiling_data
from profiling import tools
import numba


class Convolver(object):
    """Class to setup the 1D convolution of an image / mapping matrix.

    IMAGE FRAMES:
    ------------

    For a masked image, in 2D, one can compute for every unmasked pixel all other unmasked pixels it will blur light \
    into for a given PSF psf size, e.g.:

    |x|x|x|x|x|x|x|x|x|x|
    |x|x|x|x|x|x|x|x|x|x|     This is an example image.Mask, where:
    |x|x|x|x|x|x|x|x|x|x|
    |x|x|x|x|x|x|x|x|x|x|     x = True (Pixel is masked and excluded from analysis)
    |x|x|x|o|o|o|x|x|x|x|     o = False (Pixel is not masked and included in analysis)
    |x|x|x|o|o|o|x|x|x|x|
    |x|x|x|o|o|o|x|x|x|x|
    |x|x|x|x|x|x|x|x|x|x|
    |x|x|x|x|x|x|x|x|x|x|
    |x|x|x|x|x|x|x|x|x|x|

    Here, there are 9 unmasked pixels. Indexing of each unmasked pixel goes from the top-left corner right,
    and downwards, therefore:

    |x|x|x|x|x|x|x|x|x|x|
    |x|x|x|x|x|x|x|x|x|x|
    |x|x|x|x|x|x|x|x|x|x|
    |x|x|x|x|x|x|x|x|x|x|
    |x|x|x|0|1|2|x|x|x|x|
    |x|x|x|3|4|5|x|x|x|x|
    |x|x|x|6|7|8|x|x|x|x|
    |x|x|x|x|x|x|x|x|x|x|
    |x|x|x|x|x|x|x|x|x|x|
    |x|x|x|x|x|x|x|x|x|x|

    For every unmasked pixel, the Convolver over-lays the PSF psf and computes three quantities;

    image_frame_indexes - The indexes of all image pixels it will blur light into.
    image_frame_psfs - The psf values that overlap each image pixel it will blur light into.
    image_frame_length - The number of image-pixels it will blur light into (because unmasked pixels are excluded)

    For example, if we had the following 3x3 psf:

    |0.1|0.2|0.3|
    |0.4|0.5|0.6|
    |0.7|0.8|0.9|

    For pixel 0 above, when we overlap the psf 4 unmasked pixels overlap this psf, such that:

    image_frame_indexes = [0, 1, 3, 4]
    image_frame_psfs = [0.5, 0.6, 0.8, 0.9]
    image_frame_length = 4

    Noting that the other 5 psf values (0.1, 0.2, 0.3, 0.4, 0.7) overlap masked pixels and are thus discard.

    For pixel 1, we get the following results:

    image_frame_indexes = [0, 1, 2, 3, 4, 5]
    image_frame_psfs = [0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
    image_frame_lengths = 6

    In the majority of cases, there will be no unmasked pixels when the psf overlaps. This is the case above for \
    central pixel 4, where:

    image_frame_indexes = [0, 1, 2, 3, 4, 5, 6, 7, 8]
    image_frame_psfs = [0,1, 0.2, 0,3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
    image_frame_lengths = 6

    Once we have set up all these quantities, the convolution routine simply uses them to convolve a 1D array of image
    data / a mapping matrix image.

    BLURRING FRAMES:
    --------------

    Whilst the scheme above accounts for all blurred light within the mask, it does not account for the fact that
    pixels outside of the mask will also blur light into it. For galaxy light profiles, this effect is accounted for \
    using blurring frames, however it is omitted for mapping matrix images.

    First, a blurring mask is computed from a mask, which describes all pixels which are close enough to the mask \
    to blur light into it for a given psf size. Following the example above, the following blurring mask is \
    computed:

    |x|x|x|x|x|x|x|x|x|x|
    |x|x|x|x|x|x|x|x|x|x|     This is an example image.Mask, where:
    |x|x|x|x|x|x|x|x|x|x|
    |x|x|o|o|o|o|o|x|x|x|     x = True (Pixel is masked and excluded from analysis)
    |x|x|o|x|x|x|o|x|x|x|     o = False (Pixel is not masked and included in analysis)
    |x|x|o|x|x|x|o|x|x|x|
    |x|x|o|x|x|x|o|x|x|x|
    |x|x|o|o|o|o|o|x|x|x|
    |x|x|x|x|x|x|x|x|x|x|
    |x|x|x|x|x|x|x|x|x|x|

    Indexing again goes from the top-left corner right and downwards:

    |x|x| x| x| x| x| x|x|x|x|
    |x|x| x| x| x| x| x|x|x|x|
    |x|x| x| x| x| x| x|x|x|x|
    |x|x| 0| 1| 2| 3| 4|x|x|x|
    |x|x| 5| x| x| x| 6|x|x|x|
    |x|x| 7| x| x| x| 8|x|x|x|
    |x|x| 9| x| x| x|10|x|x|x|
    |x|x|11|12|13|14|15|x|x|x|
    |x|x| x| x| x| x| x|x|x|x|
    |x|x| x| x| x| x| x|x|x|x|

    For every unmasked blurring-pixel, the Convolver over-lays the PSF psf and computes three quantities;

    blurring_frame_indexes - The indexes of all unmasked image pixels it will blur light into.
    bluring_frame_kernels - The psf values that overlap each unmasked image pixel it will blur light into.
    blurring_frame_length - The number of image-pixels it will blur light into (because unmasked pixels are excluded)

    Note, therefore, that the blurring frame does not perform any blurring which blurs light into other blurring pixels.
    It only performs computations which add light inside of the mask, which is the most computationally efficient method.

    For pixel 0 above, when we overlap the 3x3 psf above only 1 unmasked image pixels overlap the psf, such that:

    blurring_frame_indexes = [0]
    blurring_frame_psfs = [0.9]
    blurring_frame_length = 1

    For pixel 1 above, when we overlap the 3x3 psf above 2 unmasked image pixels overlap the psf, such that:

    blurring_frame_indexes = [0, 1]
    blurring_frame_psfs = [0.8, 0.9]
    blurring_frame_length = 2

    For pixel 3 above, when we overlap the 3x3 psf above 3 unmasked image pixels overlap the psf, such that:

    blurring_frame_indexes = [0, 1, 2]
    blurring_frame_psfs = [0.7, 0.8, 0.9]
    blurring_frame_length = 3
    """

    mask = None
    pixels_in_mask = None
    mask_index_array = None
    kernel = None
    kernel_shape = None
    kernel_max_size = None

    def setup_mask_index_array(self):
        count = 0
        for x in range(self.mask.shape[0]):
            for y in range(self.mask.shape[1]):
                if not self.mask[x, y]:
                    self.mask_index_array[x, y] = count
                    count += 1

    def setup_image_frames(self):
        image_index = 0
        self.image_frame_indexes = np.zeros((self.pixels_in_mask, self.kernel_max_size), dtype='int')
        self.image_frame_kernels = np.zeros((self.pixels_in_mask, self.kernel_max_size))
        self.image_frame_lengths = np.zeros((self.pixels_in_mask), dtype='int')
        for x in range(self.mask_index_array.shape[0]):
            for y in range(self.mask_index_array.shape[1]):
                if not self.mask[x][y]:
                    image_frame_indexes, image_frame_kernels = \
                        self.frame_at_coords_jit((x, y), self.mask, self.mask_index_array, self.kernel,
                                                 self.kernel_shape)

                    self.image_frame_indexes[image_index, :] = image_frame_indexes
                    self.image_frame_kernels[image_index, :] = image_frame_kernels
                    self.image_frame_lengths[image_index] = image_frame_indexes[image_frame_indexes >= 0].shape[0]
                    image_index += 1

    def frame_at_coords(self, coords):
        """
        Parameters
        ----------
        coords: (int, int)
            The image_grid of mask_index_array on which the frame should be centred
        psf_shape: (int, int)
            The shape of the psf for which this frame will be used
        Returns
        -------
        frame: ndarray
            A subset of mask_index_array of shape psf_shape where elements with image_grid outside of image_frame_indexes have
            value -1
        """
        half_x = int(self.kernel_shape[0] / 2)
        half_y = int(self.kernel_shape[1] / 2)

        frame = -1 * np.ones((self.kernel_max_size))
        kernel_frame = -1.0 * np.ones((self.kernel_max_size))

        count = 0
        for i in range(self.kernel_shape[0]):
            for j in range(self.kernel_shape[1]):
                x = coords[0] - half_x + i
                y = coords[1] - half_y + j
                if 0 <= x < self.mask_index_array.shape[0] and 0 <= y < self.mask_index_array.shape[1]:
                    value = self.mask_index_array[x, y]
                    if value >= 0 and not self.mask[x, y]:
                        frame[count] = value
                        kernel_frame[count] = self.kernel[i, j]
                        count += 1

        return frame, kernel_frame

    @staticmethod
    @numba.jit(nopython=True)
    def frame_at_coords_jit(coords, mask, mask_index_array, kernel, kernel_shape):
        """
        Parameters
        ----------
        coords: (int, int)
            The image_grid of mask_index_array on which the frame should be centred
        kernel_shape: (int, int)
            The shape of the psf for which this frame will be used
        Returns
        -------
        frame: ndarray
            A subset of mask_index_array of shape psf_shape where elements with image_grid outside of image_frame_indexes have
            value -1
        """

        kernel_max_size = kernel_shape[0] * kernel_shape[1]

        half_x = int(kernel_shape[0] / 2)
        half_y = int(kernel_shape[1] / 2)

        frame = -1 * np.ones((kernel_max_size))
        kernel_frame = -1.0 * np.ones((kernel_max_size))

        count = 0
        for i in range(kernel_shape[0]):
            for j in range(kernel_shape[1]):
                x = coords[0] - half_x + i
                y = coords[1] - half_y + j
                if 0 <= x < mask_index_array.shape[0] and 0 <= y < mask_index_array.shape[1]:
                    value = mask_index_array[x, y]
                    if value >= 0 and not mask[x, y]:
                        frame[count] = value
                        kernel_frame[count] = kernel[i, j]
                        count += 1

        return frame, kernel_frame


class ConvolverImage(Convolver):

    def __init__(self, mask, blurring_mask, kernel):
        """
        Class to create image frames and blurring frames used to convolve a psf with a 1D image of non-masked \
        values.

        Parameters
        ----------
        mask : Mask
            A mask where True eliminates data.
        burring_mask : Mask
            A mask of pixels outside the mask but whose light blurs into it after convolution.
        kernel : image.PSF or ndarray
            An array representing a PSF psf.
        """

        if kernel.shape[0] % 2 == 0 or kernel.shape[1] % 2 == 0:
            raise exc.KernelException("Kernel must be odd")
        if mask.shape != blurring_mask.shape:
            raise exc.KernelException("Mask and Blurring mask must be same shape to generate Convolver")

        self.mask = mask
        self.blurring_mask = blurring_mask
        self.mask_index_array = np.full(self.mask.shape, -1)
        self.pixels_in_mask = int(np.size(self.mask) - np.sum(self.mask))
        self.pixels_in_blurring_mask = int(np.size(blurring_mask) - np.sum(blurring_mask))

        self.setup_mask_index_array()

        self.kernel = kernel
        self.kernel_shape = kernel.shape
        self.kernel_max_size = self.kernel_shape[0] * self.kernel_shape[1]

        self.setup_image_frames()
        self.setup_blurring_frames()

    def setup_blurring_frames(self):

        image_index = 0
        self.blurring_frame_indexes = np.zeros((self.pixels_in_blurring_mask, self.kernel_max_size), dtype='int')
        self.blurring_frame_kernels = np.zeros((self.pixels_in_blurring_mask, self.kernel_max_size))
        self.blurring_frame_lengths = np.zeros((self.pixels_in_blurring_mask), dtype='int')
        for x in range(self.mask.shape[0]):
            for y in range(self.mask.shape[1]):
                if self.mask[x][y] and not self.blurring_mask[x, y]:
                    image_frame_indexes, image_frame_kernels = \
                        self.frame_at_coords_jit((x, y), self.mask, self.mask_index_array, self.kernel,
                                                 self.kernel_shape)

                    self.blurring_frame_indexes[image_index, :] = image_frame_indexes
                    self.blurring_frame_kernels[image_index, :] = image_frame_kernels
                    self.blurring_frame_lengths[image_index] = image_frame_indexes[image_frame_indexes >= 0].shape[0]
                    image_index += 1

    def convolve_image(self, image_array, blurring_array):
        """
        Parameters
        ----------
        blurring_array: [Float]
            An array representing the mapping of a source pixel to a set of image pixels within the blurring region.
        sub_shape: (int, int)
            Defines a sub_grid-region of the psf for which the result should be calculated
        image_array: [float]
            A 1D array
        Returns
        -------
        convolved_vector: [float]
            A vector convolved with the psf
        """

        new_array = np.zeros(image_array.shape)

        for image_index in range(len(image_array)):

            frame_indexes = self.image_frame_indexes[image_index]
            frame_kernels = self.image_frame_kernels[image_index]
            frame_length = self.image_frame_lengths[image_index]
            value = image_array[image_index]

            for kernel_index in range(frame_length):
                vector_index = frame_indexes[kernel_index]
                kernel = frame_kernels[kernel_index]
                new_array[vector_index] += value * kernel

        for image_index in range(len(blurring_array)):

            frame_indexes = self.blurring_frame_indexes[image_index]
            frame_kernels = self.blurring_frame_kernels[image_index]
            frame_length = self.blurring_frame_lengths[image_index]
            value = blurring_array[image_index]

            for kernel_index in range(frame_length):
                vector_index = frame_indexes[kernel_index]
                kernel = frame_kernels[kernel_index]
                new_array[vector_index] += value * kernel

        return new_array


sub_grid_size=4
psf_shape = (21, 21)
# psf_shape = (41, 41)
sersic = light_profiles.EllipticalSersic(centre=(0.0, 0.0), axis_ratio=0.8, phi=90.0, intensity=0.1,
                                         effective_radius=0.8, sersic_index=4.0)

lsst = profiling_data.setup_class(name='LSST', pixel_scale=0.2, sub_grid_size=sub_grid_size, psf_shape=psf_shape)
euclid = profiling_data.setup_class(name='Euclid', pixel_scale=0.1, sub_grid_size=sub_grid_size, psf_shape=psf_shape)
hst = profiling_data.setup_class(name='HST', pixel_scale=0.05, sub_grid_size=sub_grid_size, psf_shape=psf_shape)
hst_up = profiling_data.setup_class(name='HSTup', pixel_scale=0.03, sub_grid_size=sub_grid_size, psf_shape=psf_shape)
ao = profiling_data.setup_class(name='AO', pixel_scale=0.01, sub_grid_size=sub_grid_size, psf_shape=psf_shape)

lsst_image = sersic.intensity_from_grid(grid=lsst.grids.image)
lsst_blurring_image = sersic.intensity_from_grid(grid=lsst.grids.blurring)
euclid_image = sersic.intensity_from_grid(grid=euclid.grids.image)
euclid_blurring_image = sersic.intensity_from_grid(grid=euclid.grids.blurring)
hst_image = sersic.intensity_from_grid(grid=hst.grids.image)
hst_blurring_image = sersic.intensity_from_grid(grid=hst.grids.blurring)
hst_up_image = sersic.intensity_from_grid(grid=hst_up.grids.image)
hst_up_blurring_image = sersic.intensity_from_grid(grid=hst_up.grids.blurring)
ao_image = sersic.intensity_from_grid(grid=ao.grids.image)
ao_blurring_image = sersic.intensity_from_grid(grid=ao.grids.blurring)

@tools.tick_toc_x1
def lsst_solution():
    lsst.masked_image.convolver_image.convolve_image_jit(image_array=lsst_image, blurring_array=lsst_blurring_image)

@tools.tick_toc_x1
def euclid_solution():
    euclid.masked_image.convolver_image.convolve_image_jit(image_array=euclid_image, blurring_array=euclid_blurring_image)

@tools.tick_toc_x1
def hst_solution():
    hst.masked_image.convolver_image.convolve_image_jit(image_array=hst_image, blurring_array=hst_blurring_image)

@tools.tick_toc_x1
def hst_up_solution():
    hst_up.masked_image.convolver_image.convolve_image_jit(image_array=hst_up_image, blurring_array=hst_up_blurring_image)

@tools.tick_toc_x1
def ao_solution():
    ao.masked_image.convolver_image.convolve_image_jit(image_array=ao_image, blurring_array=ao_blurring_image)

if __name__ == "__main__":
    lsst_solution()
    euclid_solution()
    hst_solution()
    hst_up_solution()
    ao_solution()