__author__ = "sibirrer"

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

from lenstronomy.Analysis.td_cosmography import TDCosmography
from lenstronomy.LensModel.lens_model import LensModel
from lenstronomy.LensModel.Solver.lens_equation_solver import LensEquationSolver
from lenstronomy.Util import constants as const


class TestTDCosmography(object):
    def setup_method(self):
        kwargs_model = {
            "lens_light_model_list": ["HERNQUIST"],
            "lens_model_list": ["SIE"],
            "point_source_model_list": ["LENSED_POSITION"],
        }
        z_lens = 0.5
        z_source = 2.5

        R_slit = 3.8
        dR_slit = 1.0
        aperture_type = "slit"
        kwargs_aperture = {
            "aperture_type": aperture_type,
            "center_ra": 0,
            "width": dR_slit,
            "length": R_slit,
            "angle": 0,
            "center_dec": 0,
        }
        psf_fwhm = 0.7
        kwargs_seeing = {"psf_type": "GAUSSIAN", "fwhm": psf_fwhm}

        TDCosmography(z_lens, z_source, kwargs_model)
        from astropy.cosmology import FlatLambdaCDM

        cosmo = FlatLambdaCDM(H0=70, Om0=0.3, Ob0=0.05)
        self.td_cosmo = TDCosmography(
            z_lens,
            z_source,
            kwargs_model,
            cosmo_fiducial=cosmo,
            lens_model_kinematics_bool=None,
            kwargs_aperture=kwargs_aperture,
            kwargs_seeing=kwargs_seeing,
            light_model_kinematics_bool=None,
        )
        self.lens = LensModel(
            lens_model_list=["SIE"], cosmo=cosmo, z_lens=z_lens, z_source=z_source
        )
        self.solver = LensEquationSolver(lensModel=self.lens)

        self.kwargs_lens = [
            {"theta_E": 1, "e1": 0.1, "e2": -0.2, "center_x": 0, "center_y": 0}
        ]
        source_x, source_y = 0, 0.05
        image_x, image_y = self.solver.image_position_from_source(
            source_x, source_y, self.kwargs_lens, min_distance=0.1, search_window=10
        )
        self.kwargs_ps = [{"ra_image": image_x, "dec_image": image_y}]
        self.image_x, self.image_y = image_x, image_y

    def test_time_delays(self):
        dt = self.td_cosmo.time_delays(self.kwargs_lens, self.kwargs_ps, kappa_ext=0)
        dt_true = self.lens.arrival_time(self.image_x, self.image_y, self.kwargs_lens)
        npt.assert_almost_equal(dt, dt_true, decimal=6)

    def test_fermat_potential(self):
        fermat_pot = self.td_cosmo.fermat_potential(self.kwargs_lens, self.kwargs_ps)
        fermat_pot_true = self.lens.fermat_potential(
            self.image_x, self.image_y, self.kwargs_lens
        )
        npt.assert_almost_equal(fermat_pot, fermat_pot_true, decimal=6)

        diff = 0.1
        kwargs_ps = [{"ra_image": self.image_x + diff, "dec_image": self.image_y}]
        fermat_pot = self.td_cosmo.fermat_potential(self.kwargs_lens, kwargs_ps)
        fermat_pot_true = self.lens.fermat_potential(
            self.image_x + diff, self.image_y, self.kwargs_lens
        )
        ratio = fermat_pot / fermat_pot_true
        assert np.max(np.abs(ratio)) > 1.05

    def test_cosmo_inference(self):
        # set up a cosmology
        # compute image postions
        # compute J and velocity dispersion
        D_dt = self.td_cosmo._lens_cosmo.ddt
        D_d = self.td_cosmo._lens_cosmo.dd
        D_s = self.td_cosmo._lens_cosmo.ds
        D_ds = self.td_cosmo._lens_cosmo.dds
        fermat_potential_list = self.td_cosmo.fermat_potential(
            self.kwargs_lens, self.kwargs_ps
        )
        dt_list = self.td_cosmo.time_delays(
            self.kwargs_lens, self.kwargs_ps, kappa_ext=0
        )
        dt = dt_list[0] - dt_list[1]
        d_fermat = fermat_potential_list[0] - fermat_potential_list[1]

        D_dt_infered = self.td_cosmo.ddt_from_time_delay(
            d_fermat_model=d_fermat, dt_measured=dt
        )
        npt.assert_almost_equal(D_dt_infered, D_dt, decimal=5)
        r_eff = 0.5
        kwargs_lens_light = [{"Rs": r_eff * 0.551, "center_x": 0, "center_y": 0}]
        kwargs_anisotropy = {"r_ani": 1}

        anisotropy_model = "OM"
        kwargs_numerics_galkin = {
            "interpol_grid_num": 500,
            "log_integration": True,
            "max_integrate": 10,
            "min_integrate": 0.001,
        }
        self.td_cosmo.kinematics_modeling_settings(
            anisotropy_model,
            kwargs_numerics_galkin,
            analytic_kinematics=True,
            Hernquist_approx=False,
            MGE_light=False,
            MGE_mass=False,
        )

        J = self.td_cosmo.velocity_dispersion_dimension_less(
            self.kwargs_lens,
            kwargs_lens_light,
            kwargs_anisotropy,
            r_eff=r_eff,
            theta_E=self.kwargs_lens[0]["theta_E"],
            gamma=2,
        )

        J_map = self.td_cosmo.velocity_dispersion_map_dimension_less(
            self.kwargs_lens,
            kwargs_lens_light,
            kwargs_anisotropy,
            r_eff=r_eff,
            theta_E=self.kwargs_lens[0]["theta_E"],
            gamma=2,
        )
        assert len(J_map) == 1
        npt.assert_almost_equal(J_map[0] / J, 1, decimal=1)
        sigma_v2 = J * D_s / D_ds * const.c**2
        sigma_v = np.sqrt(sigma_v2) / 1000.0  # convert to [km/s]
        print(sigma_v, "test sigma_v")
        Ds_Dds = self.td_cosmo.ds_dds_from_kinematics(sigma_v, J, kappa_s=0, kappa_ds=0)
        npt.assert_almost_equal(Ds_Dds, D_s / D_ds)

        # now we perform a mass-sheet transform in the observables but leave the models identical with a convergence correction
        kappa_s = 0.5
        dt_list = self.td_cosmo.time_delays(
            self.kwargs_lens, self.kwargs_ps, kappa_ext=kappa_s
        )
        sigma_v_kappa = sigma_v * np.sqrt(1 - kappa_s)
        dt = dt_list[0] - dt_list[1]
        D_dt_infered, D_d_infered = self.td_cosmo.ddt_dd_from_time_delay_and_kinematics(
            d_fermat_model=d_fermat,
            dt_measured=dt,
            sigma_v_measured=sigma_v_kappa,
            J=J,
            kappa_s=kappa_s,
            kappa_ds=0,
            kappa_d=0,
        )
        npt.assert_almost_equal(D_dt_infered, D_dt, decimal=6)
        npt.assert_almost_equal(D_d_infered, D_d, decimal=6)


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