from lenstronomy.SimulationAPI.model_api import ModelAPI
import pytest
import numpy.testing as npt
from lenstronomy.Cosmo.lens_cosmo import LensCosmo


class TestModelAPI(object):

    def setup(self):
        self.api = ModelAPI(lens_model_list=['SIS'], z_lens=None, z_source=None, lens_redshift_list=None,
                            source_light_model_list=['GAUSSIAN'], lens_light_model_list=['SERSIC'],
                            point_source_model_list=['UNLENSED'], source_redshift_list=None, cosmo=None)

    def test_lens_model_class(self):
        model = self.api.lens_model_class
        assert model.lens_model_list[0] == 'SIS'

    def test_lens_light_model_class(self):
        model = self.api.lens_light_model_class
        assert model.profile_type_list[0] == 'SERSIC'

    def test_source_model_class(self):
        model = self.api.source_model_class
        assert model.profile_type_list[0] == 'GAUSSIAN'

    def test_point_source_model_class(self):
        model = self.api.point_source_model_class
        assert model.point_source_type_list[0] == 'UNLENSED'

    def test_source_position(self):
        api = ModelAPI(lens_model_list=['SIS'], z_lens=None, z_source=None, lens_redshift_list=None,
                            source_light_model_list=['GAUSSIAN'], lens_light_model_list=['SERSIC'],
                            point_source_model_list=['SOURCE_POSITION'], source_redshift_list=None, cosmo=None)
        model = api.point_source_model_class
        assert model.point_source_type_list[0] == 'SOURCE_POSITION'

    def test_physical2lensing_conversion(self):
        lens_redshift_list = [0.5, 1]
        z_source_convention = 2
        api = ModelAPI(lens_model_list=['SIS', 'NFW'], lens_redshift_list=lens_redshift_list,
                       z_source_convention=z_source_convention, cosmo=None, z_source=z_source_convention)

        kwargs_mass = [{'sigma_v': 200, 'center_x': 0, 'center_y': 0},
                       {'M200': 10**13, 'concentration': 5, 'center_x': 1, 'center_y': 1}]
        kwargs_lens = api.physical2lensing_conversion(kwargs_mass)

        theta_E = kwargs_lens[0]['theta_E']
        lens_cosmo = LensCosmo(z_lens=lens_redshift_list[0], z_source=z_source_convention)
        theta_E_test = lens_cosmo.sis_sigma_v2theta_E(kwargs_mass[0]['sigma_v'])
        npt.assert_almost_equal(theta_E, theta_E_test, decimal=7)

        alpha_Rs = kwargs_lens[1]['alpha_Rs']
        lens_cosmo = LensCosmo(z_lens=lens_redshift_list[1], z_source=z_source_convention)
        Rs_new , alpha_Rs_new = lens_cosmo.nfw_physical2angle(kwargs_mass[1]['M200'], kwargs_mass[1]['concentration'])
        npt.assert_almost_equal(alpha_Rs, alpha_Rs_new, decimal=7)

        api = ModelAPI(lens_model_list=['SIS', 'NFW'], z_lens=0.5, z_source_convention=z_source_convention, cosmo=None,
                       z_source=z_source_convention)
        kwargs_lens = api.physical2lensing_conversion(kwargs_mass)

        theta_E = kwargs_lens[0]['theta_E']
        lens_cosmo = LensCosmo(z_lens=0.5, z_source=z_source_convention)
        theta_E_test = lens_cosmo.sis_sigma_v2theta_E(kwargs_mass[0]['sigma_v'])
        npt.assert_almost_equal(theta_E, theta_E_test, decimal=7)


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