__author__ = "sibirrer"

import numpy.testing as npt
import numpy as np
from lenstronomy.LensModel.lens_model_extensions import LensModelExtensions
from lenstronomy.LensModel.Solver.lens_equation_solver import LensEquationSolver
from lenstronomy.LensModel.lens_model import LensModel
from lenstronomy.Util.magnification_finite_util import (
    auto_raytracing_grid_resolution,
    auto_raytracing_grid_size,
)
import lenstronomy.Util.param_util as param_util
from lenstronomy.LightModel.light_model import LightModel
from astropy.cosmology import FlatLambdaCDM


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

    def setup_method(self):
        self.cosmo = FlatLambdaCDM(H0=70, Om0=0.3)

    def test_critical_curves(self):
        lens_model_list = ["SPEP"]
        phi, q = 1.0, 0.8
        e1, e2 = param_util.phi_q2_ellipticity(phi, q)
        kwargs_lens = [
            {
                "theta_E": 1.0,
                "gamma": 2.0,
                "e1": e1,
                "e2": e2,
                "center_x": 0,
                "center_y": 0,
            }
        ]
        lens_model = LensModel(lens_model_list)
        lensModelExtensions = LensModelExtensions(LensModel(lens_model_list))

        # this first test is for even pixel grid numbers
        compute_window, grid_scale = 5, 0.05
        num_pix = int(compute_window / grid_scale)
        assert num_pix % 2 == 0
        (
            ra_crit_list,
            dec_crit_list,
            ra_caustic_list,
            dec_caustic_list,
        ) = lensModelExtensions.critical_curve_caustics(
            kwargs_lens, compute_window=compute_window, grid_scale=grid_scale
        )

        # here we test whether the caustic points are in fact at high magnifications (close to infinite)
        # close here means above magnification of 1000000 (with matplotlib method, this limit achieved was 170)
        for k in range(len(ra_crit_list)):
            ra_crit = ra_crit_list[k]
            dec_crit = dec_crit_list[k]
            mag = lens_model.magnification(ra_crit, dec_crit, kwargs_lens)
            assert np.all(np.abs(mag) > 1000)

        # this second test is for odd pixel grid numbers
        compute_window, grid_scale = 5.0 + 0.06, 0.05
        print(compute_window / grid_scale, "test float")
        num_pix = int(compute_window / grid_scale)
        print(num_pix, "test num_pix")
        assert num_pix % 2 == 1
        (
            ra_crit_list,
            dec_crit_list,
            ra_caustic_list,
            dec_caustic_list,
        ) = lensModelExtensions.critical_curve_caustics(
            kwargs_lens, compute_window=compute_window, grid_scale=grid_scale
        )

        # here we test whether the caustic points are in fact at high magnifications (close to infinite)
        # close here means above magnification of 1000000 (with matplotlib method, this limit achieved was 170)
        for k in range(len(ra_crit_list)):
            ra_crit = ra_crit_list[k]
            dec_crit = dec_crit_list[k]
            mag = lens_model.magnification(ra_crit, dec_crit, kwargs_lens)
            assert np.all(np.abs(mag) > 1000)

    def test_critical_curves_tiling(self):
        lens_model_list = ["SPEP"]
        phi, q = 1.0, 0.8
        e1, e2 = param_util.phi_q2_ellipticity(phi, q)
        kwargs_lens = [
            {
                "theta_E": 1.0,
                "gamma": 2.0,
                "e1": e1,
                "e2": e2,
                "center_x": 0,
                "center_y": 0,
            }
        ]
        lensModel = LensModelExtensions(LensModel(lens_model_list))
        ra_crit, dec_crit = lensModel.critical_curve_tiling(
            kwargs_lens, compute_window=5, start_scale=0.01, max_order=10
        )
        # here we test whether the caustic points are in fact at high magnifications (close to infinite)
        # close here means above magnification of 1000. This is more precise than the critical_curve_caustics() method
        lens_model = LensModel(lens_model_list)
        mag = lens_model.magnification(ra_crit, dec_crit, kwargs_lens)
        assert np.all(np.abs(mag) > 1000)

    def test_get_magnification_model(self):
        self.kwargs_options = {
            "lens_model_list": ["GAUSSIAN"],
            "source_light_model_list": ["GAUSSIAN"],
            "lens_light_model_list": ["SERSIC"],
            "subgrid_res": 10,
            "numPix": 200,
            "psf_type": "gaussian",
            "x2_simple": True,
        }
        kwargs_lens = [
            {"amp": 1, "sigma_x": 2, "sigma_y": 2, "center_x": 0, "center_y": 0}
        ]

        x_pos = np.array([1.0, 1.0, 2.0])
        y_pos = np.array([-1.0, 0.0, 0.0])
        lens_model = LensModelExtensions(LensModel(lens_model_list=["GAUSSIAN"]))
        mag = lens_model.magnification_finite(
            x_pos,
            y_pos,
            kwargs_lens,
            source_sigma=0.003,
            window_size=0.1,
            grid_number=100,
        )
        npt.assert_almost_equal(mag[0], 0.98848384784633392, decimal=5)

    def test_magnification_finite(self):
        lens_model_list = ["SPEP", "SHEAR"]

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

        extension = LensModelExtensions(LensModel(lens_model_list))
        x_image = [0.56153533, -0.78067875, -0.72551184, 0.75664112]
        y_image = [-0.74722528, 0.52491177, -0.72799235, 0.78503659]

        mag_square_grid = extension.magnification_finite(
            x_image,
            y_image,
            kwargs_lens,
            source_sigma=0.001,
            grid_number=200,
            window_size=0.1,
        )
        mag_polar_grid = extension.magnification_finite(
            x_image,
            y_image,
            kwargs_lens,
            source_sigma=0.001,
            grid_number=200,
            window_size=0.1,
            polar_grid=True,
        )
        npt.assert_almost_equal(mag_polar_grid, mag_square_grid, decimal=5)

    def test_magnification_finite_adaptive(self):
        lens_model_list = ["EPL", "SHEAR"]
        z_source = 1.5
        kwargs_lens = [
            {
                "theta_E": 1.0,
                "gamma": 2.0,
                "e1": 0.02,
                "e2": -0.09,
                "center_x": 0,
                "center_y": 0,
            },
            {"gamma1": 0.01, "gamma2": 0.03},
        ]

        lensmodel = LensModel(lens_model_list)
        extension = LensModelExtensions(lensmodel)
        solver = LensEquationSolver(lensmodel)
        source_x, source_y = 0.07, 0.03
        x_image, y_image = solver.findBrightImage(source_x, source_y, kwargs_lens)

        source_fwhm_parsec = 40.0

        pc_per_arcsec = 1000 / self.cosmo.arcsec_per_kpc_proper(z_source).value
        source_sigma = source_fwhm_parsec / pc_per_arcsec / 2.355

        grid_size = auto_raytracing_grid_size(source_fwhm_parsec)
        grid_resolution = auto_raytracing_grid_resolution(source_fwhm_parsec)
        # make this even higher resolution to show convergence
        grid_number = int(2 * grid_size / grid_resolution)
        window_size = 2 * grid_size

        mag_square_grid = extension.magnification_finite(
            x_image,
            y_image,
            kwargs_lens,
            source_sigma=source_sigma,
            grid_number=grid_number,
            window_size=window_size,
        )
        flux_ratios_square_grid = mag_square_grid / max(mag_square_grid)

        mag_adaptive_grid = extension.magnification_finite_adaptive(
            x_image,
            y_image,
            source_x,
            source_y,
            kwargs_lens,
            source_fwhm_parsec,
            z_source,
            cosmo=self.cosmo,
        )
        flux_ratios_adaptive_grid = mag_adaptive_grid / max(mag_adaptive_grid)

        mag_adaptive_grid_fixed_aperture_size = extension.magnification_finite_adaptive(
            x_image,
            y_image,
            source_x,
            source_y,
            kwargs_lens,
            source_fwhm_parsec,
            z_source,
            cosmo=self.cosmo,
            fixed_aperture_size=True,
            grid_radius_arcsec=0.2,
        )
        flux_ratios_fixed_aperture_size = mag_adaptive_grid_fixed_aperture_size / max(
            mag_adaptive_grid_fixed_aperture_size
        )

        mag_adaptive_grid_2 = extension.magnification_finite_adaptive(
            x_image,
            y_image,
            source_x,
            source_y,
            kwargs_lens,
            source_fwhm_parsec,
            z_source,
            cosmo=self.cosmo,
            axis_ratio=0,
        )
        mag_adaptive_grid_3 = extension.magnification_finite_adaptive(
            x_image,
            y_image,
            source_x,
            source_y,
            kwargs_lens,
            source_fwhm_parsec,
            z_source,
            cosmo=self.cosmo,
            axis_ratio=1,
        )

        flux_ratios_adaptive_grid_2 = mag_adaptive_grid_2 / max(mag_adaptive_grid_2)
        flux_ratios_adaptive_grid_3 = mag_adaptive_grid_3 / max(mag_adaptive_grid_3)

        # tests the default cosmology
        _ = extension.magnification_finite_adaptive(
            x_image,
            y_image,
            source_x,
            source_y,
            kwargs_lens,
            source_fwhm_parsec,
            z_source,
            cosmo=None,
            tol=0.0001,
        )

        # test smallest eigenvalue
        _ = extension.magnification_finite_adaptive(
            x_image,
            y_image,
            source_x,
            source_y,
            kwargs_lens,
            source_fwhm_parsec,
            z_source,
            cosmo=None,
            tol=0.0001,
            use_largest_eigenvalue=False,
        )

        # tests the r_max > sqrt(2) * grid_radius stop criterion
        _ = extension.magnification_finite_adaptive(
            x_image,
            y_image,
            source_x,
            source_y,
            kwargs_lens,
            source_fwhm_parsec,
            z_source,
            cosmo=None,
            tol=0.0001,
            step_size=1000,
        )

        mag_point_source = abs(lensmodel.magnification(x_image, y_image, kwargs_lens))

        quarter_precent_precision = [0.0025] * 4
        npt.assert_array_less(
            flux_ratios_square_grid / flux_ratios_adaptive_grid - 1,
            quarter_precent_precision,
        )
        npt.assert_array_less(
            flux_ratios_fixed_aperture_size / flux_ratios_adaptive_grid - 1,
            quarter_precent_precision,
        )
        npt.assert_array_less(
            flux_ratios_square_grid / flux_ratios_adaptive_grid_2 - 1,
            quarter_precent_precision,
        )
        npt.assert_array_less(
            flux_ratios_adaptive_grid_3 / flux_ratios_adaptive_grid_2 - 1,
            quarter_precent_precision,
        )
        half_percent_precision = [0.005] * 4
        npt.assert_array_less(
            mag_square_grid / mag_adaptive_grid - 1, half_percent_precision
        )
        npt.assert_array_less(
            mag_square_grid / mag_adaptive_grid_2 - 1, half_percent_precision
        )
        npt.assert_array_less(
            mag_adaptive_grid / mag_point_source - 1, half_percent_precision
        )

        flux_array = np.array([0.0, 0.0])
        grid_x = np.array([0.0, source_sigma])
        grid_y = np.array([0.0, 0.0])
        grid_r = np.hypot(grid_x, grid_y)

        # test that the double gaussian profile has 2x the flux when dx, dy = 0
        magnification_double_gaussian = extension.magnification_finite_adaptive(
            x_image,
            y_image,
            source_x,
            source_y,
            kwargs_lens,
            source_fwhm_parsec,
            z_source,
            cosmo=self.cosmo,
            source_light_model="DOUBLE_GAUSSIAN",
            dx=0.0,
            dy=0.0,
            amp_scale=1.0,
            size_scale=1.0,
        )
        npt.assert_almost_equal(magnification_double_gaussian, 2 * mag_adaptive_grid)

        grid_radius = 0.3
        npix = 400
        _x = _y = np.linspace(-grid_radius, grid_radius, npix)
        resolution = 2 * grid_radius / npix
        xx, yy = np.meshgrid(_x, _y)
        for i in range(0, 4):
            beta_x, beta_y = lensmodel.ray_shooting(
                x_image[i] + xx.ravel(), y_image[i] + yy.ravel(), kwargs_lens
            )
            source_light_model = LightModel(["GAUSSIAN"] * 2)
            amp_scale, dx, dy, size_scale = 0.2, 0.005, -0.005, 1.0
            kwargs_source = [
                {
                    "amp": 1.0,
                    "center_x": source_x,
                    "center_y": source_y,
                    "sigma": source_sigma,
                },
                {
                    "amp": amp_scale,
                    "center_x": source_x + dx,
                    "center_y": source_y + dy,
                    "sigma": source_sigma * size_scale,
                },
            ]
            sb = source_light_model.surface_brightness(beta_x, beta_y, kwargs_source)
            magnification_double_gaussian_reference = np.sum(sb) * resolution**2
            magnification_double_gaussian = extension.magnification_finite_adaptive(
                [x_image[i]],
                [y_image[i]],
                source_x,
                source_y,
                kwargs_lens,
                source_fwhm_parsec,
                z_source,
                cosmo=self.cosmo,
                source_light_model="DOUBLE_GAUSSIAN",
                dx=dx,
                dy=dy,
                amp_scale=amp_scale,
                size_scale=size_scale,
                grid_resolution=resolution,
                grid_radius_arcsec=grid_radius,
                axis_ratio=1.0,
            )
            npt.assert_almost_equal(
                magnification_double_gaussian / magnification_double_gaussian_reference,
                1.0,
                3,
            )

        source_model = LightModel(["GAUSSIAN"])
        kwargs_source = [
            {
                "amp": 1.0,
                "center_x": source_x,
                "center_y": source_y,
                "sigma": source_sigma,
            }
        ]

        r_min = 0.0
        r_max = source_sigma * 0.9
        flux_array = extension._magnification_adaptive_iteration(
            flux_array,
            x_image[0],
            y_image[0],
            grid_x,
            grid_y,
            grid_r,
            r_min,
            r_max,
            lensmodel,
            kwargs_lens,
            source_model,
            kwargs_source,
        )
        bx, by = lensmodel.ray_shooting(x_image[0], y_image[0], kwargs_lens)
        sb_true = source_model.surface_brightness(bx, by, kwargs_source)
        npt.assert_equal(True, flux_array[0] == sb_true)
        npt.assert_equal(True, flux_array[1] == 0.0)

        r_min = source_sigma * 0.9
        r_max = 2 * source_sigma

        flux_array = extension._magnification_adaptive_iteration(
            flux_array,
            x_image[0],
            y_image[0],
            grid_x,
            grid_y,
            grid_r,
            r_min,
            r_max,
            lensmodel,
            kwargs_lens,
            source_model,
            kwargs_source,
        )
        bx, by = lensmodel.ray_shooting(
            x_image[0] + source_sigma, y_image[0], kwargs_lens
        )
        sb_true = source_model.surface_brightness(bx, by, kwargs_source)
        npt.assert_equal(True, flux_array[1] == sb_true)

    def test_zoom_source(self):
        lens_model_list = ["SIE", "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, "e1": 0.1, "e2": 0.1, "center_x": 0, "center_y": 0},
            {"gamma1": 0.05, "gamma2": -0.03},
        ]

        x_img, y_img = lensEquationSolver.image_position_from_source(
            kwargs_lens=kwargs_lens, sourcePos_x=x_source, sourcePos_y=y_source
        )

        image = lensModelExtensions.zoom_source(
            x_img[0],
            y_img[0],
            kwargs_lens,
            source_sigma=0.003,
            window_size=0.1,
            grid_number=100,
            shape="GAUSSIAN",
        )
        assert len(image) == 100

    def test_tangential_average(self):
        lens_model_list = ["SIS"]
        lensModel = LensModel(lens_model_list=lens_model_list)
        lensModelExtensions = LensModelExtensions(lensModel=lensModel)
        tang_stretch_ave = lensModelExtensions.tangential_average(
            x=1.1,
            y=0,
            kwargs_lens=[{"theta_E": 1, "center_x": 0, "center_y": 0}],
            dr=1,
            smoothing=None,
            num_average=9,
        )
        npt.assert_almost_equal(tang_stretch_ave, -2.525501464097973, decimal=6)

    def test_caustic_area(self):
        lens_model_list = ["SIE"]
        lensModel = LensModel(lens_model_list=lens_model_list)
        lensModelExtensions = LensModelExtensions(lensModel=lensModel)
        kwargs_lens = [{"theta_E": 1, "e1": 0.2, "e2": 0, "center_x": 0, "center_y": 0}]
        kwargs_caustic_num = {
            "compute_window": 3,
            "grid_scale": 0.005,
            "center_x": 0,
            "center_y": 0,
        }

        area = lensModelExtensions.caustic_area(
            kwargs_lens=kwargs_lens,
            kwargs_caustic_num=kwargs_caustic_num,
            index_vertices=0,
        )
        npt.assert_almost_equal(area, 0.08445866728739478, decimal=3)
