from lenstronomy.PointSource.Types.source_position import SourcePositions
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 = SourcePositions(lens_model=lens_model, fixed_magnification=True)
        self.ps = SourcePositions(lens_model=lens_model, fixed_magnification=False)
        self.kwargs = {'point_amp': [2, 1], 'ra_source': 0.1, 'dec_source': 0}
        self.kwargs_mag = {'source_amp': 1, 'ra_source': 0.1, 'dec_source': 0}

    def test_image_position(self):
        x_img, y_img = self.ps.image_position(self.kwargs, self.kwargs_lens)

        lens_model = LensModel(lens_model_list=['SIS'])
        x_src, y_src = lens_model.ray_shooting(x_img, y_img, kwargs=self.kwargs_lens)
        npt.assert_almost_equal(x_src, self.kwargs['ra_source'])
        npt.assert_almost_equal(y_src, self.kwargs['dec_source'])

    def test_source_position(self):
        x_src, y_src = self.ps.source_position(self.kwargs, kwargs_lens=self.kwargs_lens)
        npt.assert_almost_equal(x_src, self.kwargs['ra_source'])
        npt.assert_almost_equal(y_src, self.kwargs['dec_source'])

    def test_image_amplitude(self):
        x_img, y_img = self.ps.image_position(self.kwargs, kwargs_lens=self.kwargs_lens)
        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=x_img,
                                              y_pos=y_img, magnification_limit=None, kwargs_lens_eqn_solver=None)
        npt.assert_almost_equal(amp_pos, amp)

        amp = self.ps.image_amplitude(self.kwargs, kwargs_lens=self.kwargs_lens, x_pos=x_img,
                                              y_pos=y_img, magnification_limit=None,
                                              kwargs_lens_eqn_solver=None)
        npt.assert_almost_equal(amp, self.kwargs['point_amp'])

        #see if works with mag_pert defined
        self.kwargs['mag_pert'] = [0.1,0.1]
        amp_pert = self.ps.image_amplitude(self.kwargs, kwargs_lens=self.kwargs_lens, x_pos=x_img,
                                              y_pos=y_img, magnification_limit=None,
                                              kwargs_lens_eqn_solver=None)
        npt.assert_almost_equal(amp_pert, 0.1*amp)

    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()
