__author__ = "sibirrer"

import numpy as np
import numpy.testing as npt
from lenstronomy.ImSim.Numerics.convolution import (
    MultiGaussianConvolution,
    PixelKernelConvolution,
    SubgridKernelConvolution,
    MGEConvolution,
)
from lenstronomy.LightModel.light_model import LightModel
import lenstronomy.Util.util as util
import pytest


class TestPixelKernelConvolution(object):
    def setup_method(self):
        lightModel = LightModel(light_model_list=["GAUSSIAN"])
        self.delta_pix = 1
        x, y = util.make_grid(10, deltapix=self.delta_pix)
        kwargs = [{"amp": 1, "sigma": 1, "center_x": 0, "center_y": 0}]
        flux = lightModel.surface_brightness(x, y, kwargs)
        self.model = util.array2image(flux)

    def test_convolve2d(self):
        kernel = np.zeros((3, 3))
        kernel[1, 1] = 1
        pixel_conv = PixelKernelConvolution(kernel=kernel)
        image_convolved = pixel_conv.convolution2d(self.model)
        npt.assert_almost_equal(np.sum(image_convolved), np.sum(self.model), decimal=2)

    def test_copy_transpose(self):
        kernel = np.zeros((3, 3))
        kernel[1, 1] = 1
        pixel_conv = PixelKernelConvolution(kernel=kernel)
        pixel_conv_t = pixel_conv.copy_transpose()
        image_convolved = pixel_conv.convolution2d(self.model)
        image_convolved_t = pixel_conv_t.convolution2d(self.model)
        npt.assert_equal(image_convolved, image_convolved_t)

    def test_pixel_kernel(self):
        kernel = np.zeros((5, 5))
        kernel[1, 1] = 1
        pixel_conv = PixelKernelConvolution(kernel=kernel)
        npt.assert_equal(pixel_conv.pixel_kernel(), kernel)
        npt.assert_equal(pixel_conv.pixel_kernel(num_pix=3), kernel[1:-1, 1:-1])


class TestSubgridKernelConvolution(object):
    def setup_method(self):
        self.supersampling_factor = 3
        lightModel = LightModel(light_model_list=["GAUSSIAN"])
        self.delta_pix = 1.0
        x, y = util.make_grid(20, deltapix=self.delta_pix)
        x_sub, y_sub = util.make_grid(
            20 * self.supersampling_factor,
            deltapix=self.delta_pix / self.supersampling_factor,
        )
        kwargs = [{"amp": 1, "sigma": 2, "center_x": 0, "center_y": 0}]
        flux = lightModel.surface_brightness(x, y, kwargs)
        self.model = util.array2image(flux)
        flux_sub = lightModel.surface_brightness(x_sub, y_sub, kwargs)
        self.model_sub = util.array2image(flux_sub)

        x, y = util.make_grid(5, deltapix=self.delta_pix)
        kwargs_kernel = [{"amp": 1, "sigma": 1, "center_x": 0, "center_y": 0}]
        kernel = lightModel.surface_brightness(x, y, kwargs_kernel)
        self.kernel = util.array2image(kernel) / np.sum(kernel)

        x_sub, y_sub = util.make_grid(
            5 * self.supersampling_factor,
            deltapix=self.delta_pix / self.supersampling_factor,
        )
        kernel_sub = lightModel.surface_brightness(x_sub, y_sub, kwargs_kernel)
        self.kernel_sub = util.array2image(kernel_sub) / np.sum(kernel_sub)

    def test_fft_scipy_static(self):
        supersampling_factor = 2
        conv_sicpy = SubgridKernelConvolution(
            self.kernel,
            supersampling_factor,
            supersampling_kernel_size=None,
            convolution_type="fft",
        )

        conv_static = SubgridKernelConvolution(
            self.kernel,
            supersampling_factor,
            supersampling_kernel_size=None,
            convolution_type="fft_static",
        )

        model_conv_scipy = conv_sicpy.convolution2d(self.model)
        model_conv_static = conv_static.convolution2d(self.model)
        npt.assert_almost_equal(model_conv_static, model_conv_scipy, decimal=3)

    def test_convolve2d(self):
        # kernel_supersampled = kernel_util.subgrid_kernel(self.kernel, self.supersampling_factor, odd=True, num_iter=5)
        subgrid_conv = SubgridKernelConvolution(
            self.kernel_sub,
            self.supersampling_factor,
            supersampling_kernel_size=None,
            convolution_type="fft",
        )
        model_subgrid_conv = subgrid_conv.convolution2d(self.model_sub)

        supersampling_factor = 1
        conv = SubgridKernelConvolution(
            self.kernel,
            supersampling_factor,
            supersampling_kernel_size=None,
            convolution_type="fft",
        )
        model_conv = conv.convolution2d(self.model)
        npt.assert_almost_equal(
            np.sum(model_subgrid_conv), np.sum(model_conv), decimal=1
        )
        npt.assert_almost_equal(model_subgrid_conv, model_conv, decimal=1)

        # kernel_supersampled = kernel_util.subgrid_kernel(self.kernel, self.supersampling_factor, odd=True, num_iter=5)
        subgrid_conv_split = SubgridKernelConvolution(
            self.kernel_sub,
            self.supersampling_factor,
            supersampling_kernel_size=5,
            convolution_type="fft",
        )
        model_subgrid_conv_split = subgrid_conv_split.convolution2d(self.model_sub)
        npt.assert_almost_equal(
            np.sum(model_subgrid_conv), np.sum(model_subgrid_conv_split), decimal=8
        )
        npt.assert_almost_equal(model_subgrid_conv, model_subgrid_conv_split, decimal=8)

        subgrid_conv_split = SubgridKernelConvolution(
            self.kernel_sub,
            self.supersampling_factor,
            supersampling_kernel_size=3,
            convolution_type="fft",
        )
        model_subgrid_conv_split = subgrid_conv_split.convolution2d(self.model_sub)
        npt.assert_almost_equal(
            np.sum(model_subgrid_conv), np.sum(model_subgrid_conv_split), decimal=5
        )
        npt.assert_almost_equal(model_subgrid_conv, model_subgrid_conv_split, decimal=3)


class TestMultiGaussianConvolution(object):
    def setup_method(self):
        lightModel = LightModel(light_model_list=["GAUSSIAN"])
        self.delta_pix = 1
        x, y = util.make_grid(10, deltapix=self.delta_pix)
        kwargs = [{"amp": 1, "sigma": 1, "center_x": 0, "center_y": 0}]
        flux = lightModel.surface_brightness(x, y, kwargs)
        self.model = util.array2image(flux)

    def test_convolve2d(self):
        sigma_list = [0.5, 1, 2]
        fraction_list = [0.5, 0.2, 0.3]
        mge_conv = MultiGaussianConvolution(
            sigma_list=sigma_list,
            fraction_list=fraction_list,
            pixel_scale=self.delta_pix,
        )
        image_convolved = mge_conv.convolution2d(self.model)
        npt.assert_almost_equal(np.sum(image_convolved), np.sum(self.model), decimal=2)


class TestMGEConvolution(object):
    def setup_method(self):
        lightModel = LightModel(light_model_list=["GAUSSIAN"])
        self.delta_pix = 1
        x, y = util.make_grid(10, deltapix=self.delta_pix)
        kwargs = [{"amp": 1, "sigma": 2, "center_x": 0, "center_y": 0}]
        flux = lightModel.surface_brightness(x, y, kwargs)
        self.model = util.array2image(flux)

    def test_convolve2d(self):
        sigma_list = [2, 3, 4]
        fraction_list = [0.5, 0.2, 0.3]
        mg_conv = MultiGaussianConvolution(
            sigma_list=sigma_list,
            fraction_list=fraction_list,
            pixel_scale=self.delta_pix,
        )
        pixel_kernel = mg_conv.pixel_kernel(num_pix=11)
        mge_conv = MGEConvolution(pixel_kernel, pixel_scale=self.delta_pix, order=20)
        image_conv_mg = mg_conv.convolution2d(self.model)
        image_conv_mge = mge_conv.convolution2d(self.model)
        npt.assert_almost_equal(
            image_conv_mge / np.max(image_conv_mg),
            image_conv_mg / np.max(image_conv_mg),
            decimal=2,
        )

        diff_kernel = mge_conv.kernel_difference()
        npt.assert_almost_equal(
            diff_kernel,
            pixel_kernel - mge_conv._mge_conv.pixel_kernel(len(pixel_kernel)),
        )


if __name__ == "__main__":
    pytest.main()
