__author__ = "sibirrer"

import numpy.testing as npt
import numpy as np
import pytest

import lenstronomy.Util.param_util as param_util
from lenstronomy.LensModel.lens_model import LensModel
from lenstronomy.LightModel.light_model import LightModel
from lenstronomy.PointSource.point_source import PointSource
from lenstronomy.ImSim.image_model import ImageModel
from lenstronomy.ImSim.image_linear_solve import ImageLinearFit
import lenstronomy.Util.simulation_util as sim_util
from lenstronomy.LensModel.Solver.lens_equation_solver import LensEquationSolver
from lenstronomy.Data.imaging_data import ImageData
from lenstronomy.Data.psf import PSF
from lenstronomy.ImSim.differential_extinction import DifferentialExtinction
from lenstronomy.Util import util


class TestImageModel(object):
    """Tests the source model routines."""

    def setup_method(self):
        # data specifics
        sigma_bkg = 0.05  # background noise per pixel
        exp_time = 100  # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit)
        numPix = 100  # cutout pixel size
        deltaPix = 0.05  # pixel size in arcsec (area per pixel = deltaPix**2)
        fwhm = 0.5  # full width half max of PSF

        # PSF specification

        kwargs_data = sim_util.data_configure_simple(
            numPix, deltaPix, exp_time, sigma_bkg, inverse=True
        )
        data_class = ImageData(**kwargs_data)
        kwargs_psf = {
            "psf_type": "GAUSSIAN",
            "fwhm": fwhm,
            "truncation": 5,
            "pixel_size": deltaPix,
        }
        psf_class = PSF(**kwargs_psf)
        kernel = psf_class.kernel_point_source
        kwargs_psf = {
            "psf_type": "PIXEL",
            "kernel_point_source": kernel,
            "psf_error_map": np.ones_like(kernel) * 0.001 * kernel**2,
        }
        psf_class = PSF(**kwargs_psf)

        # 'EXERNAL_SHEAR': external shear
        kwargs_shear = {
            "gamma1": 0.01,
            "gamma2": 0.01,
        }  # gamma_ext: shear strength, psi_ext: shear angel (in radian)
        phi, q = 0.2, 0.8
        e1, e2 = param_util.phi_q2_ellipticity(phi, q)
        kwargs_spemd = {
            "theta_E": 1.0,
            "gamma": 1.8,
            "center_x": 0,
            "center_y": 0,
            "e1": e1,
            "e2": e2,
        }

        lens_model_list = ["SPEP", "SHEAR"]
        self.kwargs_lens = [kwargs_spemd, kwargs_shear]
        lens_model_class = LensModel(lens_model_list=lens_model_list)
        # list of light profiles (for lens and source)
        # 'SERSIC': spherical Sersic profile
        kwargs_sersic = {
            "amp": 1.0,
            "R_sersic": 0.1,
            "n_sersic": 2,
            "center_x": 0,
            "center_y": 0,
        }
        # 'SERSIC_ELLIPSE': elliptical Sersic profile
        phi, q = 0.2, 0.9
        e1, e2 = param_util.phi_q2_ellipticity(phi, q)
        kwargs_sersic_ellipse = {
            "amp": 1.0,
            "R_sersic": 0.6,
            "n_sersic": 7,
            "center_x": 0,
            "center_y": 0,
            "e1": e1,
            "e2": e2,
        }

        lens_light_model_list = ["SERSIC"]
        self.kwargs_lens_light = [kwargs_sersic]
        lens_light_model_class = LightModel(light_model_list=lens_light_model_list)
        source_model_list = ["SERSIC_ELLIPSE"]
        self.kwargs_source = [kwargs_sersic_ellipse]
        source_model_class = LightModel(light_model_list=source_model_list)
        self.kwargs_ps = [
            {"ra_source": 0.01, "dec_source": 0.0, "source_amp": 1.0}
        ]  # quasar point source position in the source plane and intrinsic brightness
        point_source_class = PointSource(
            point_source_type_list=["SOURCE_POSITION"], fixed_magnification_list=[True]
        )
        kwargs_numerics = {
            "supersampling_factor": 2,
            "supersampling_convolution": False,
        }
        imageModel = ImageModel(
            data_class,
            psf_class,
            lens_model_class,
            source_model_class,
            lens_light_model_class,
            point_source_class,
            kwargs_numerics=kwargs_numerics,
        )
        image_sim = sim_util.simulate_simple(
            imageModel,
            self.kwargs_lens,
            self.kwargs_source,
            self.kwargs_lens_light,
            self.kwargs_ps,
        )
        data_class.update_data(image_sim)

        self.imageModel = ImageLinearFit(
            data_class,
            psf_class,
            lens_model_class,
            source_model_class,
            lens_light_model_class,
            point_source_class,
            kwargs_numerics=kwargs_numerics,
        )
        self.solver = LensEquationSolver(lensModel=self.imageModel.LensModel)

    def test_source_surface_brightness(self):
        source_model = self.imageModel.source_surface_brightness(
            self.kwargs_source, self.kwargs_lens, unconvolved=False, de_lensed=True
        )
        assert len(source_model) == 100

        source_model = self.imageModel.source_surface_brightness(
            self.kwargs_source, self.kwargs_lens, unconvolved=False, de_lensed=False
        )
        assert len(source_model) == 100
        npt.assert_almost_equal(
            source_model[10, 10], 0.13939841209844345 * 0.05**2, decimal=4
        )

        source_model = self.imageModel.source_surface_brightness(
            self.kwargs_source, self.kwargs_lens, unconvolved=True, de_lensed=False
        )
        assert len(source_model) == 100
        npt.assert_almost_equal(
            source_model[10, 10], 0.13536114618182182 * 0.05**2, decimal=4
        )

    def test_lens_surface_brightness(self):
        lens_flux = self.imageModel.lens_surface_brightness(
            self.kwargs_lens_light, unconvolved=False
        )
        print(np.sum(lens_flux), "test lens flux")
        npt.assert_almost_equal(lens_flux[50, 50], 0.0010788981265391802, decimal=4)
        # npt.assert_almost_equal(lens_flux[50, 50], 0.54214440654021534 * 0.05 ** 2, decimal=4)

        lens_flux = self.imageModel.lens_surface_brightness(
            self.kwargs_lens_light, unconvolved=True
        )
        npt.assert_almost_equal(
            lens_flux[50, 50], 4.7310552067454452 * 0.05**2, decimal=4
        )

    def test_image_linear_solve(self):
        model, error_map, cov_param, param = self.imageModel.image_linear_solve(
            self.kwargs_lens,
            self.kwargs_source,
            self.kwargs_lens_light,
            self.kwargs_ps,
            inv_bool=False,
        )
        chi2_reduced = self.imageModel.reduced_chi2(model, error_map)
        npt.assert_almost_equal(chi2_reduced, 1, decimal=1)

    def test_linear_response_matrix(self):
        A = self.imageModel.linear_response_matrix(
            self.kwargs_lens, self.kwargs_source, self.kwargs_lens_light, self.kwargs_ps
        )
        n, m = np.shape(A)
        assert n == 3
        assert m == 100 * 100

    def test_image_with_params(self):
        model = self.imageModel.image(
            self.kwargs_lens,
            self.kwargs_source,
            self.kwargs_lens_light,
            self.kwargs_ps,
            unconvolved=False,
            source_add=True,
            lens_light_add=True,
            point_source_add=True,
        )
        error_map = self.imageModel._error_map_psf(self.kwargs_lens, self.kwargs_ps)
        chi2_reduced = self.imageModel.reduced_chi2(model, error_map)
        npt.assert_almost_equal(chi2_reduced, 1, decimal=1)

    def test_likelihood_data_given_model(self):
        logL = self.imageModel.likelihood_data_given_model(
            self.kwargs_lens,
            self.kwargs_source,
            self.kwargs_lens_light,
            self.kwargs_ps,
            source_marg=False,
        )
        npt.assert_almost_equal(logL, -5000, decimal=-3)

        logLmarg = self.imageModel.likelihood_data_given_model(
            self.kwargs_lens,
            self.kwargs_source,
            self.kwargs_lens_light,
            self.kwargs_ps,
            source_marg=True,
        )
        npt.assert_almost_equal(logL - logLmarg, 0, decimal=-3)
        assert logLmarg < logL

    def test_reduced_residuals(self):
        model = sim_util.simulate_simple(
            self.imageModel,
            self.kwargs_lens,
            self.kwargs_source,
            self.kwargs_lens_light,
            self.kwargs_ps,
            no_noise=True,
        )
        residuals = self.imageModel.reduced_residuals(model, error_map=0)
        npt.assert_almost_equal(np.std(residuals), 1.01, decimal=1)

        chi2 = self.imageModel.reduced_chi2(model, error_map=0)
        npt.assert_almost_equal(chi2, 1, decimal=1)

    def test_numData_evaluate(self):
        numData = self.imageModel.num_data_evaluate
        assert numData == 10000

    def test_num_param_linear(self):
        num_param_linear = self.imageModel.num_param_linear(
            self.kwargs_lens, self.kwargs_source, self.kwargs_lens_light, self.kwargs_ps
        )
        assert num_param_linear == 3

    def test_update_data(self):
        kwargs_data = sim_util.data_configure_simple(
            numPix=10, deltaPix=1, exposure_time=1, background_rms=1, inverse=True
        )
        data_class = ImageData(**kwargs_data)
        self.imageModel.update_data(data_class)
        assert self.imageModel.Data.num_pixel == 100

    def test_point_source_rendering(self):
        # initialize data

        numPix = 100
        deltaPix = 0.05
        kwargs_data = sim_util.data_configure_simple(
            numPix, deltaPix, exposure_time=1, background_rms=1
        )
        data_class = ImageData(**kwargs_data)
        kernel = np.zeros((5, 5))
        kernel[2, 2] = 1
        kwargs_psf = {
            "kernel_point_source": kernel,
            "psf_type": "PIXEL",
            "psf_error_map": np.ones_like(kernel) * 0.001,
        }
        psf_class = PSF(**kwargs_psf)
        lens_model_class = LensModel(["SPEP"])
        source_model_class = LightModel([])
        lens_light_model_class = LightModel([])
        kwargs_numerics = {
            "supersampling_factor": 2,
            "supersampling_convolution": True,
            "point_source_supersampling_factor": 1,
        }
        point_source_class = PointSource(
            point_source_type_list=["LENSED_POSITION"], fixed_magnification_list=[False]
        )
        makeImage = ImageModel(
            data_class,
            psf_class,
            lens_model_class,
            source_model_class,
            lens_light_model_class,
            point_source_class,
            kwargs_numerics=kwargs_numerics,
        )
        # chose point source positions
        x_pix = np.array([10, 5, 10, 90])
        y_pix = np.array([40, 50, 60, 50])
        ra_pos, dec_pos = makeImage.Data.map_pix2coord(x_pix, y_pix)
        e1, e2 = param_util.phi_q2_ellipticity(0, 0.8)
        kwargs_lens_init = [
            {"theta_E": 1, "gamma": 2, "e1": e1, "e2": e2, "center_x": 0, "center_y": 0}
        ]
        kwargs_else = [
            {
                "ra_image": ra_pos,
                "dec_image": dec_pos,
                "point_amp": np.ones_like(ra_pos),
            }
        ]
        image = makeImage.image(
            kwargs_lens_init,
            kwargs_source={},
            kwargs_lens_light={},
            kwargs_ps=kwargs_else,
        )
        # print(np.shape(model), 'test')
        # image = makeImage.ImageNumerics.array2image(model)
        for i in range(len(x_pix)):
            npt.assert_almost_equal(image[y_pix[i], x_pix[i]], 1, decimal=2)

        x_pix = np.array([10.5, 5.5, 10.5, 90.5])
        y_pix = np.array([40, 50, 60, 50])
        ra_pos, dec_pos = makeImage.Data.map_pix2coord(x_pix, y_pix)
        phi, q = 0.0, 0.8
        e1, e2 = param_util.phi_q2_ellipticity(phi, q)
        kwargs_lens_init = [
            {"theta_E": 1, "gamma": 2, "e1": e1, "e2": e2, "center_x": 0, "center_y": 0}
        ]
        kwargs_else = [
            {
                "ra_image": ra_pos,
                "dec_image": dec_pos,
                "point_amp": np.ones_like(ra_pos),
            }
        ]
        image = makeImage.image(
            kwargs_lens_init,
            kwargs_source={},
            kwargs_lens_light={},
            kwargs_ps=kwargs_else,
        )
        # image = makeImage.ImageNumerics.array2image(model)
        for i in range(len(x_pix)):
            print(int(y_pix[i]), int(x_pix[i] + 0.5))
            npt.assert_almost_equal(image[int(y_pix[i]), int(x_pix[i])], 0.5, decimal=1)
            npt.assert_almost_equal(
                image[int(y_pix[i]), int(x_pix[i] + 0.5)], 0.5, decimal=1
            )

    def test_point_source(self):
        pointSource = PointSource(
            point_source_type_list=["SOURCE_POSITION"], fixed_magnification_list=[True]
        )
        kwargs_ps = [{"source_amp": 1000, "ra_source": 0.1, "dec_source": 0.1}]
        lensModel = LensModel(lens_model_list=["SIS"])
        kwargs_lens = [{"theta_E": 1, "center_x": 0, "center_y": 0}]
        numPix = 64
        deltaPix = 0.13
        kwargs_data = sim_util.data_configure_simple(
            numPix, deltaPix, exposure_time=1, background_rms=1
        )
        data_class = ImageData(**kwargs_data)

        psf_type = "GAUSSIAN"
        fwhm = 0.9
        kwargs_psf = {"psf_type": psf_type, "fwhm": fwhm}
        psf_class = PSF(**kwargs_psf)
        imageModel = ImageModel(
            data_class=data_class,
            psf_class=psf_class,
            lens_model_class=lensModel,
            point_source_class=pointSource,
        )
        image = imageModel.image(kwargs_lens=kwargs_lens, kwargs_ps=kwargs_ps)
        assert np.sum(image) > 0

    def test_error_map_source(self):
        sourceModel = LightModel(light_model_list=["UNIFORM", "UNIFORM"])

        kwargs_data = sim_util.data_configure_simple(
            numPix=10, deltaPix=1, exposure_time=1, background_rms=1
        )
        data_class = ImageData(**kwargs_data)

        psf_type = "GAUSSIAN"
        fwhm = 0.9
        kwargs_psf = {"psf_type": psf_type, "fwhm": fwhm}
        psf_class = PSF(**kwargs_psf)
        imageModel = ImageLinearFit(
            data_class=data_class,
            psf_class=psf_class,
            lens_model_class=None,
            point_source_class=None,
            source_model_class=sourceModel,
        )

        x_grid, y_grid = util.make_grid(numPix=10, deltapix=1)
        error_map = imageModel.error_map_source(
            kwargs_source=[{"amp": 1}, {"amp": 1}],
            x_grid=x_grid,
            y_grid=y_grid,
            cov_param=np.array([[1, 0], [0, 1]]),
        )
        assert error_map[0] == 2

    def test_create_empty(self):
        kwargs_data = sim_util.data_configure_simple(
            numPix=10, deltaPix=1, exposure_time=1, background_rms=1
        )
        data_class = ImageData(**kwargs_data)
        imageModel_empty = ImageModel(data_class, PSF())
        assert imageModel_empty._psf_error_map == False

        flux = imageModel_empty.lens_surface_brightness(kwargs_lens_light=None)
        assert flux.all() == 0

    def test_extinction_map(self):
        kwargs_data = sim_util.data_configure_simple(
            numPix=10, deltaPix=1, exposure_time=1, background_rms=1
        )
        data_class = ImageData(**kwargs_data)
        extinction_class = DifferentialExtinction(
            optical_depth_model=["UNIFORM"], tau0_index=0
        )
        imageModel = ImageModel(data_class, PSF(), extinction_class=extinction_class)
        extinction = imageModel.extinction_map(
            kwargs_extinction=[{"amp": 1}], kwargs_special={"tau0_list": [1, 0, 0]}
        )
        npt.assert_almost_equal(extinction, np.exp(-1))

    def test_error_response(self):
        C_D_response, model_error = self.imageModel._error_response(
            self.kwargs_lens, self.kwargs_ps, kwargs_special=None
        )
        assert len(model_error) == 100
        print(np.sum(model_error))
        npt.assert_almost_equal(np.sum(model_error), 0.0019271126921470687, decimal=3)

    def test_point_source_linear_response_set(self):
        kwargs_special = {"delta_x_image": [0.1, 0.1], "delta_y_image": [-0.1, -0.1]}
        (
            ra_pos,
            dec_pos,
            amp,
            num_point,
        ) = self.imageModel.point_source_linear_response_set(
            self.kwargs_ps, self.kwargs_lens, kwargs_special, with_amp=True
        )
        ra, dec = self.imageModel.PointSource.image_position(
            self.kwargs_ps, self.kwargs_lens
        )
        npt.assert_almost_equal(ra[0][0], ra_pos[0][0] - 0.1, decimal=5)

    def test_displace_astrometry(self):
        kwargs_special = {
            "delta_x_image": np.array([0.1, 0.1]),
            "delta_y_image": np.array([-0.1, -0.1]),
        }
        x_pos, y_pos = np.array([0, 0]), np.array([0, 0])
        x_shift, y_shift = self.imageModel._displace_astrometry(
            x_pos, y_pos, kwargs_special=kwargs_special
        )
        assert x_pos[0] == 0
        assert x_shift[0] == kwargs_special["delta_x_image"][0]
        assert y_pos[0] == 0
        assert y_shift[0] == kwargs_special["delta_y_image"][0]


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