from lenstronomy.ImSim.Numerics.point_source_rendering import PointSourceRendering
from lenstronomy.Data.pixel_grid import PixelGrid
from lenstronomy.Data.psf import PSF

import numpy as np
import numpy.testing as npt
import pytest
import unittest


class TestPointSourceRendering(object):
    def setup_method(self):
        Mpix2coord = np.array([[1, 0], [0, 1]])
        kwargs_grid = {
            "ra_at_xy_0": 0,
            "dec_at_xy_0": 0,
            "transform_pix2angle": Mpix2coord,
            "nx": 10,
            "ny": 10,
        }
        pixel_grid = PixelGrid(**kwargs_grid)
        kernel = np.zeros((5, 5))
        kernel[2, 2] = 1
        kwargs_psf = {
            "kernel_point_source": kernel,
            "psf_type": "PIXEL",
            "psf_variance_map": np.ones_like(kernel) * kernel**2,
        }
        psf_class = PSF(**kwargs_psf)

        self._ps_rendering = PointSourceRendering(
            pixel_grid, supersampling_factor=1, psf=psf_class
        )

    def test_psf_variance_map(self):
        ra_pos, dec_pos = [5], [5]
        data = np.zeros((10, 10))
        image = self._ps_rendering.psf_variance_map(
            ra_pos, dec_pos, amp=1, data=data, fix_psf_variance_map=False
        )
        npt.assert_almost_equal(np.sum(image), 0, decimal=10)

        image = self._ps_rendering.psf_variance_map(
            ra_pos, dec_pos, amp=1, data=data, fix_psf_variance_map=True
        )
        npt.assert_almost_equal(np.sum(image), 1, decimal=10)

        ra_pos, dec_pos = [50], [50]
        data = np.zeros((10, 10))
        image = self._ps_rendering.psf_variance_map(
            ra_pos, dec_pos, amp=1, data=data, fix_psf_variance_map=False
        )
        npt.assert_almost_equal(np.sum(image), 0, decimal=10)

    def test_point_source_rendering(self):
        amp = [1, 1]
        ra_pos, dec_pos = [0, 1], [1, 0]
        model = self._ps_rendering.point_source_rendering(ra_pos, dec_pos, amp)
        npt.assert_almost_equal(np.sum(model), 2, decimal=8)


class TestRaise(unittest.TestCase):
    def test_raise(self):
        Mpix2coord = np.array([[1, 0], [0, 1]])
        kwargs_grid = {
            "ra_at_xy_0": 0,
            "dec_at_xy_0": 0,
            "transform_pix2angle": Mpix2coord,
            "nx": 10,
            "ny": 10,
        }
        pixel_grid = PixelGrid(**kwargs_grid)
        kernel = np.zeros((5, 5))
        kernel[2, 2] = 1
        kwargs_psf = {
            "kernel_point_source": kernel,
            "psf_type": "PIXEL",
            "psf_variance_map": np.ones_like(kernel),
        }
        psf_class = PSF(**kwargs_psf)

        self._ps_rendering = PointSourceRendering(
            pixel_grid, supersampling_factor=1, psf=psf_class
        )
        with self.assertRaises(ValueError):
            self._ps_rendering.point_source_rendering(
                ra_pos=[1, 1], dec_pos=[0, 1], amp=[1]
            )


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