import pytest
import numpy as np
import numpy.testing as npt
from lenstronomy.Sampling.Likelihoods.flux_ratio_likelihood import FluxRatioLikelihood
from lenstronomy.LensModel.lens_model import LensModel
from lenstronomy.LensModel.lens_model_extensions import LensModelExtensions
from lenstronomy.LensModel.Solver.lens_equation_solver import LensEquationSolver


class TestFluxRatioLikelihood(object):
    def setup_method(self):
        lens_model_list = ["SPEP", "SHEAR"]
        lensModel = LensModel(lens_model_list=lens_model_list)
        lensModelExtensions = LensModelExtensions(lensModel=lensModel)
        lensEquationSolver = LensEquationSolver(lensModel=lensModel)

        x_source, y_source = 0.02, 0.01
        kwargs_lens = [
            {
                "theta_E": 1.0,
                "e1": 0.1,
                "e2": 0.1,
                "gamma": 2.0,
                "center_x": 0,
                "center_y": 0,
            },
            {"gamma1": 0.06, "gamma2": -0.03},
        ]

        x_img, y_img = lensEquationSolver.image_position_from_source(
            kwargs_lens=kwargs_lens, sourcePos_x=x_source, sourcePos_y=y_source
        )
        print("image positions are: ", x_img, y_img)
        mag_inf = lensModel.magnification(x_img, y_img, kwargs_lens)
        print("point source magnification: ", mag_inf)

        source_size_arcsec = 0.001
        window_size = 0.1
        grid_number = 100
        print("source size in arcsec: ", source_size_arcsec)
        mag_finite = lensModelExtensions.magnification_finite(
            x_pos=x_img,
            y_pos=y_img,
            kwargs_lens=kwargs_lens,
            source_sigma=source_size_arcsec,
            window_size=window_size,
            grid_number=grid_number,
        )
        flux_ratios = mag_finite[1:] / mag_finite[0]
        flux_ratio_errors = [0.1, 0.1, 0.1]
        flux_ratio_cov = np.diag([0.1, 0.1, 0.1]) ** 2
        self.flux_likelihood = FluxRatioLikelihood(
            lens_model_class=lensModel,
            flux_ratios=flux_ratios,
            flux_ratio_errors=flux_ratio_errors,
            source_type="GAUSSIAN",
            window_size=window_size,
            grid_number=grid_number,
        )

        self.flux_likelihood_inf = FluxRatioLikelihood(
            lens_model_class=lensModel,
            flux_ratios=flux_ratios,
            flux_ratio_errors=flux_ratio_errors,
            source_type="INF",
            window_size=window_size,
            grid_number=grid_number,
        )
        self.flux_likelihood_inf_cov = FluxRatioLikelihood(
            lens_model_class=lensModel,
            flux_ratios=flux_ratios,
            flux_ratio_errors=flux_ratio_cov,
            source_type="INF",
            window_size=window_size,
            grid_number=grid_number,
        )
        self.kwargs_cosmo = {"source_size": source_size_arcsec}
        self.x_img, self.y_img = x_img, y_img
        self.kwargs_lens = kwargs_lens

    def test_logL(self):
        logL = self.flux_likelihood.logL(
            self.x_img, self.y_img, self.kwargs_lens, kwargs_special=self.kwargs_cosmo
        )
        assert logL == 0

        logL_inf = self.flux_likelihood_inf.logL(
            self.x_img, self.y_img, self.kwargs_lens, {}
        )
        npt.assert_almost_equal(logL_inf, 0, decimal=4)

    def test__logL(self):
        lensModel = LensModel(lens_model_list=[])
        flux_ratios_init = np.array([1.0, 1.0, 1.0])
        flux_ratio_errors = np.array([1.0, 1.0, 1.0])
        flux_likelihood = FluxRatioLikelihood(
            lens_model_class=lensModel,
            flux_ratios=flux_ratios_init,
            flux_ratio_errors=flux_ratio_errors,
        )

        flux_ratios = np.array([0, 1, np.nan])
        logL = flux_likelihood._logL(flux_ratios)
        assert logL == -(10**15)

        flux_likelihood = FluxRatioLikelihood(
            lens_model_class=lensModel,
            flux_ratios=flux_ratios_init,
            flux_ratio_errors=np.array([0.0, 1.0, 1.0]),
        )
        flux_ratios = np.array([1.0, 1.0, 1.0])
        logL = flux_likelihood._logL(flux_ratios)
        assert logL == -(10**15)

    def test_numimgs(self):
        # Test with a different number of images
        logL = self.flux_likelihood.logL(
            self.x_img[:-1],
            self.y_img[:-1],
            self.kwargs_lens,
            kwargs_special=self.kwargs_cosmo,
        )
        assert logL == -(10**15)

    def test_covmatrix(self):
        # Test with a different number of images
        logL = self.flux_likelihood_inf_cov.logL(
            self.x_img, self.y_img, self.kwargs_lens, kwargs_special=self.kwargs_cosmo
        )
        npt.assert_almost_equal(logL, 0, decimal=8)


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