__author__ = "sibirrer"

import numpy as np
import numpy.testing as npt
from lenstronomy.ImSim.Numerics.adaptive_numerics import AdaptiveConvolution
from lenstronomy.ImSim.Numerics.convolution import SubgridKernelConvolution
from lenstronomy.LightModel.light_model import LightModel
import lenstronomy.Util.util as util
import pytest


class TestAdaptiveConvolution(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_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)

        conv_supersample_pixels = np.zeros_like(self.model)
        conv_supersample_pixels = np.array(conv_supersample_pixels, dtype=bool)
        conv_supersample_pixels[self.model > np.max(self.model) / 20] = True
        adaptive_conv = AdaptiveConvolution(
            self.kernel_sub,
            self.supersampling_factor,
            conv_supersample_pixels,
            supersampling_kernel_size=5,
            compute_pixels=None,
        )

        model_adaptive_conv = adaptive_conv.convolve2d(self.model_sub)
        npt.assert_almost_equal(
            np.sum(model_subgrid_conv), np.sum(model_adaptive_conv), decimal=2
        )
        npt.assert_almost_equal(model_subgrid_conv, model_adaptive_conv, decimal=2)

        conv_supersample_pixels = np.zeros_like(self.model)
        conv_supersample_pixels = np.array(conv_supersample_pixels, dtype=bool)
        conv_supersample_pixels[self.model > np.max(self.model) / 2] = True
        adaptive_conv = AdaptiveConvolution(
            self.kernel_sub,
            self.supersampling_factor,
            conv_supersample_pixels,
            supersampling_kernel_size=1,
            compute_pixels=None,
        )

        model_adaptive_conv = adaptive_conv.convolve2d(self.model_sub)

        npt.assert_almost_equal(
            np.sum(model_subgrid_conv), np.sum(model_adaptive_conv), decimal=2
        )
        npt.assert_almost_equal(model_subgrid_conv, model_adaptive_conv, decimal=2)


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