from lenstronomy.ImSim.differential_extinction import DifferentialExtinction
import pytest
import numpy.testing as npt


class TestDifferentialExtinction(object):
    def setup_method(self):
        pass

    def test_extinction(self):
        extinction = DifferentialExtinction(
            optical_depth_model=["GAUSSIAN"], tau0_index=0
        )
        kwargs_extinction = [{"amp": 1, "sigma": 1, "center_x": 0, "center_y": 0}]
        kwargs_special = {"tau0_list": [2, 0]}
        ext = extinction.extinction(
            x=1, y=1, kwargs_special=kwargs_special, kwargs_extinction=kwargs_extinction
        )
        npt.assert_almost_equal(ext, 0.8894965388088921, decimal=8)

        ext = extinction.extinction(
            x=1, y=1, kwargs_special=kwargs_special, kwargs_extinction=None
        )
        npt.assert_almost_equal(ext, 1, decimal=8)

        ext = extinction.extinction(
            x=1, y=1, kwargs_special={}, kwargs_extinction=kwargs_extinction
        )
        npt.assert_almost_equal(ext, 0.9431312415612645, decimal=8)


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