from lenstronomy.SimulationAPI.sim_api import SimAPI
import lenstronomy.SimulationAPI.observation_constructor as constructor
import pytest
import numpy.testing as npt


class TestModelAPI(object):
    def setup_method(self):
        numpix = 10
        instrument_name = "LSST"
        observation_name = "LSST_g_band"
        kwargs_single_band = constructor.observation_constructor(
            instrument_name=instrument_name, observation_name=observation_name
        )
        kwargs_single_band["data_count_unit"] = "e-"
        kwargs_model = {
            "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,
        }
        kwargs_numerics = {"supersampling_factor": 2}
        self.api = SimAPI(numpix, kwargs_single_band, kwargs_model)

    def test_image_model_class(self):
        model = self.api.image_model_class()
        assert model.LensModel.lens_model_list[0] == "SIS"

    def test_magnitude2amplitude(self):
        kwargs_lens_light_mag = [
            {
                "magnitude": 28,
                "R_sersic": 1.0,
                "n_sersic": 2,
                "center_x": 0,
                "center_y": 0,
            }
        ]
        kwargs_source_mag = [
            {"magnitude": 30, "sigma": 0.3, "center_x": 0, "center_y": 0}
        ]
        kwargs_ps_mag = [{"magnitude": [30], "ra_image": [0], "dec_image": [0]}]
        kwargs_lens_light, kwargs_source, kwargs_ps = self.api.magnitude2amplitude(
            kwargs_lens_light_mag, kwargs_source_mag, kwargs_ps_mag
        )

        npt.assert_almost_equal(kwargs_source[0]["amp"], 1, decimal=5)
        npt.assert_almost_equal(kwargs_ps[0]["point_amp"][0], 1, decimal=5)
        npt.assert_almost_equal(
            kwargs_lens_light[0]["amp"], 0.38680586575451237, decimal=5
        )


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