from lenstronomy.PointSource.Types.lensed_position import LensedPositions
from lenstronomy.LensModel.lens_model import LensModel
import pytest
import numpy.testing as npt


class TestLensedPosition(object):

    def setup_method(self):
        lens_model = LensModel(lens_model_list=['SIS'])
        self.kwargs_lens = [{'theta_E': 1, 'center_x': 0, 'center_y': 0}]
        self.ps_mag = LensedPositions(lens_model=lens_model, fixed_magnification=True, point_source_frame_list=[0, 0],
                                      index_lens_model_list=[[0]])
        self.ps = LensedPositions(lens_model=lens_model, fixed_magnification=False)
        self.ps_add = LensedPositions(lens_model=lens_model, fixed_magnification=[False], additional_images=True)
        self.kwargs = {'point_amp': [2, 1], 'ra_image': [0, 1.2], 'dec_image': [0, 0]}
        self.kwargs_mag = {'source_amp': 2, 'ra_image': [0, 1.2], 'dec_image': [0, 0]}

    def test_image_source_position(self):
        x_img, y_img = self.ps.image_position(self.kwargs, self.kwargs_lens)
        npt.assert_almost_equal(x_img, self.kwargs['ra_image'])

        x_img_add, y_img_add = self.ps_add.image_position(self.kwargs, self.kwargs_lens)
        print(x_img_add, x_img)
        assert x_img[0] != x_img_add[0]

        # check whether the source solution matches
        x_src, y_src = self.ps.source_position(self.kwargs, self.kwargs_lens)
        lens_model = LensModel(lens_model_list=['SIS'])
        x_src_true, y_src_true = lens_model.ray_shooting(x_img_add, y_img_add, kwargs=self.kwargs_lens)
        npt.assert_almost_equal(x_src_true[0], x_src_true[1])
        npt.assert_almost_equal(y_src_true[0], y_src_true[1])

        npt.assert_almost_equal(x_src_true, x_src)
        npt.assert_almost_equal(y_src_true, y_src)

    def test_image_amplitude(self):
        amp = self.ps.image_amplitude(self.kwargs, kwargs_lens=self.kwargs_lens, x_pos=self.kwargs['ra_image'],
                                      y_pos=self.kwargs['dec_image'], magnification_limit=None,
                                      kwargs_lens_eqn_solver=None)
        npt.assert_almost_equal(self.kwargs['point_amp'], amp)

        amp = self.ps_mag.image_amplitude(self.kwargs_mag, kwargs_lens=self.kwargs_lens, x_pos=None,
                                          y_pos=None, magnification_limit=None, kwargs_lens_eqn_solver=None)

        amp_pos = self.ps_mag.image_amplitude(self.kwargs_mag, kwargs_lens=self.kwargs_lens, x_pos=self.kwargs['ra_image'],
                                              y_pos=self.kwargs['dec_image'], magnification_limit=None,
                                              kwargs_lens_eqn_solver=None)
        npt.assert_almost_equal(amp, amp_pos)

    def test_source_amplitude(self):
        amp = self.ps.source_amplitude(self.kwargs, kwargs_lens=self.kwargs_lens)
        amp_mag = self.ps_mag.source_amplitude(self.kwargs_mag, kwargs_lens=self.kwargs_lens)
        npt.assert_almost_equal(amp_mag, self.kwargs_mag['source_amp'])
        assert amp != amp_mag


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