__author__ = "mgomer"

from lenstronomy.LensModel.Profiles.synthesis import SynthesisProfile
from lenstronomy.Analysis.lens_profile import LensProfileAnalysis
from lenstronomy.LensModel.lens_model import LensModel

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


class TestSynthesis(object):
    """Tests the synthesis model's ability to approximate several profiles."""

    def setup_method(self):
        self.lin_fit_hyperparams = {"num_r_evals": 100}
        self.s_list = np.logspace(-6.0, 3.0, 30)
        self.x_test = np.linspace(0.01, 2, 10)
        self.y_test = np.zeros_like(self.x_test)

    def test_CSE_components(self):
        # test potential, deflection, and kappa using CSE components for a few profiles
        kwargs_list = []
        center_x = 0
        center_y = 0
        for s in self.s_list:
            kwargs_list.append(
                {
                    "a": 1,
                    "s": s,
                    "e1": 0,
                    "e2": 0,
                    "center_x": center_x,
                    "center_y": center_y,
                }
            )

        # test nfw
        kwargs_nfw = [
            {"Rs": 1.5, "alpha_Rs": 1, "center_x": center_x, "center_y": center_y}
        ]
        kwargs_synthesis = {
            "target_lens_model": "NFW",
            "component_lens_model": "CSE",
            "kwargs_list": kwargs_list,
            "lin_fit_hyperparams": self.lin_fit_hyperparams,
        }
        lensmodel_synth = LensModel(["SYNTHESIS"], kwargs_synthesis=kwargs_synthesis)
        lensmodel_nfw = LensModel(["NFW"])
        self.compare_synth(lensmodel_synth, lensmodel_nfw, kwargs_nfw)
        # test sersic
        kwargs_sersic = [{"k_eff": 3, "R_sersic": 1.5, "n_sersic": 3.5}]
        kwargs_synthesis = {
            "target_lens_model": "SERSIC",
            "component_lens_model": "CSE",
            "kwargs_list": kwargs_list,
            "lin_fit_hyperparams": self.lin_fit_hyperparams,
        }
        lensmodel_synth = LensModel(["SYNTHESIS"], kwargs_synthesis=kwargs_synthesis)
        lensmodel_sersic = LensModel(["SERSIC"])
        self.compare_synth(lensmodel_synth, lensmodel_sersic, kwargs_sersic)
        # test hernquist
        kwargs_hernquist = [
            {"sigma0": 10, "Rs": 1.5, "center_x": center_x, "center_y": center_y}
        ]
        kwargs_synthesis = {
            "target_lens_model": "HERNQUIST",
            "component_lens_model": "CSE",
            "kwargs_list": kwargs_list,
            "lin_fit_hyperparams": self.lin_fit_hyperparams,
        }
        lensmodel_synth = LensModel(["SYNTHESIS"], kwargs_synthesis=kwargs_synthesis)
        lensmodel_hernquist = LensModel(["HERNQUIST"])
        self.compare_synth(lensmodel_synth, lensmodel_hernquist, kwargs_hernquist)

    def test_gaussian_components(self):
        # test potential, deflection, and kappa using Gaussian components for a few profiles
        kwargs_list = []
        center_x = 0
        center_y = 0
        for s in self.s_list:
            kwargs_list.append(
                {"amp": 1, "sigma": s, "e1": 0, "e2": 0, "center_x": 0, "center_y": 0}
            )
        # test nfw
        kwargs_nfw = [
            {"Rs": 1.5, "alpha_Rs": 1, "center_x": center_x, "center_y": center_y}
        ]
        kwargs_synthesis = {
            "target_lens_model": "NFW",
            "component_lens_model": "GAUSSIAN_ELLIPSE_KAPPA",
            "kwargs_list": kwargs_list,
            "lin_fit_hyperparams": self.lin_fit_hyperparams,
        }
        lensmodel_synth = LensModel(["SYNTHESIS"], kwargs_synthesis=kwargs_synthesis)
        lensmodel_nfw = LensModel(["NFW"])
        self.compare_synth(lensmodel_synth, lensmodel_nfw, kwargs_nfw)
        # test sersic
        kwargs_sersic = [{"k_eff": 3, "R_sersic": 1.5, "n_sersic": 3.5}]
        kwargs_synthesis = {
            "target_lens_model": "SERSIC",
            "component_lens_model": "GAUSSIAN_ELLIPSE_KAPPA",
            "kwargs_list": kwargs_list,
            "lin_fit_hyperparams": self.lin_fit_hyperparams,
        }
        lensmodel_synth = LensModel(["SYNTHESIS"], kwargs_synthesis=kwargs_synthesis)
        lensmodel_sersic = LensModel(["SERSIC"])
        self.compare_synth(lensmodel_synth, lensmodel_sersic, kwargs_sersic)
        # test hernquist
        kwargs_hernquist = [
            {"sigma0": 10, "Rs": 1.5, "center_x": center_x, "center_y": center_y}
        ]
        kwargs_synthesis = {
            "target_lens_model": "HERNQUIST",
            "component_lens_model": "GAUSSIAN_ELLIPSE_KAPPA",
            "kwargs_list": kwargs_list,
            "lin_fit_hyperparams": self.lin_fit_hyperparams,
        }
        lensmodel_synth = LensModel(["SYNTHESIS"], kwargs_synthesis=kwargs_synthesis)
        lensmodel_hernquist = LensModel(["HERNQUIST"])
        self.compare_synth(lensmodel_synth, lensmodel_hernquist, kwargs_hernquist)

    def compare_synth(self, lensmodel_synth, lensmodel_target, kwargs_target):
        # check potentials, deflection, kappa
        synth_pot = lensmodel_synth.potential(self.x_test, self.y_test, kwargs_target)
        target_pot = lensmodel_target.potential(self.x_test, self.y_test, kwargs_target)
        synth_defl = lensmodel_synth.alpha(self.x_test, self.y_test, kwargs_target)[0]
        target_defl = lensmodel_target.alpha(self.x_test, self.y_test, kwargs_target)[0]
        synth_kappa = lensmodel_synth.kappa(self.x_test, self.y_test, kwargs_target)
        target_kappa = lensmodel_target.kappa(self.x_test, self.y_test, kwargs_target)

        npt.assert_allclose(synth_pot, target_pot, rtol=1e-2)  # potentials within 1%
        npt.assert_allclose(synth_defl, target_defl, rtol=1e-2)  # deflections within 1%
        npt.assert_allclose(synth_kappa, target_kappa, rtol=1e-2)  # kappa within 1%

    def test_ellipticity_and_centers(self):
        # check that even with ellipticity and offset center, same radial profile
        r_test = self.x_test
        kwargs_list = []
        center_x = 1
        center_y = -1
        for s in self.s_list:
            kwargs_list.append(
                {
                    "a": 1,
                    "s": s,
                    "e1": 0.3,
                    "e2": 0,
                    "center_x": center_x,
                    "center_y": center_y,
                }
            )

        # test nfw from CSEs
        kwargs_nfw = [
            {"Rs": 1.5, "alpha_Rs": 1, "center_x": center_x, "center_y": center_y}
        ]
        kwargs_synthesis = {
            "target_lens_model": "NFW",
            "component_lens_model": "CSE",
            "kwargs_list": kwargs_list,
            "lin_fit_hyperparams": self.lin_fit_hyperparams,
        }
        lensmodel_synth = LensModel(["SYNTHESIS"], kwargs_synthesis=kwargs_synthesis)
        lensmodel_nfw = LensModel(["NFW"])
        LensAn_synth = LensProfileAnalysis(lensmodel_synth)
        synth_avg_kappa = LensAn_synth.radial_lens_profile(r_test, kwargs_nfw)
        LensAn_nfw = LensProfileAnalysis(lensmodel_nfw)
        nfw_avg_kappa = LensAn_nfw.radial_lens_profile(r_test, kwargs_nfw)
        npt.assert_allclose(nfw_avg_kappa, synth_avg_kappa, rtol=5e-2)

    def test_set_static_method(self):
        # test that when set_static is used the result doesn't load the kwargs anymore, instead using the saved weights
        kwargs_list = []
        center_x = 0
        center_y = 0
        for s in self.s_list:
            kwargs_list.append(
                {
                    "a": 1,
                    "s": s,
                    "e1": 0,
                    "e2": 0,
                    "center_x": center_x,
                    "center_y": center_y,
                }
            )

        # test nfw
        kwargs_nfw = [
            {"Rs": 1.5, "alpha_Rs": 1, "center_x": center_x, "center_y": center_y}
        ]
        kwargs_synthesis = {
            "target_lens_model": "NFW",
            "component_lens_model": "CSE",
            "kwargs_list": kwargs_list,
            "lin_fit_hyperparams": self.lin_fit_hyperparams,
        }
        synth_prof = SynthesisProfile(**kwargs_synthesis)
        synth_func = synth_prof.function(self.x_test, self.y_test, **kwargs_nfw[0])
        linear_weights = synth_prof.linear_weight_mle_fit(kwargs_nfw, kwargs_list)
        synth_prof.set_static(linear_weights)
        kwargs_nfw[0]["Rs"] = 10  # change kwargs, doesn't affect static weights
        changed_synth_func = synth_prof.function(
            self.x_test, self.y_test, **kwargs_nfw[0]
        )
        npt.assert_array_equal(synth_func, changed_synth_func)
