__author__ = "sibirrer"

import lenstronomy.Util.multi_gauss_expansion as mge

import numpy as np
import numpy.testing as npt
from lenstronomy.LightModel.Profiles.sersic import Sersic
from lenstronomy.LightModel.Profiles.hernquist import Hernquist
from lenstronomy.LightModel.Profiles.gaussian import MultiGaussian
import pytest


class TestMGE(object):
    """Tests the Gaussian methods."""

    def setup_method(self):
        self.sersic = Sersic()
        self.multiGaussian = MultiGaussian()

    def test_mge_1d_sersic(self):
        n_comp = 30
        r_sersic = 1.0
        n_sersic = 3.7
        I0_sersic = 1.0
        rs = np.logspace(-2.0, 1.0, 50) * r_sersic
        ss = self.sersic.function(
            rs, np.zeros_like(rs), amp=I0_sersic, n_sersic=n_sersic, R_sersic=r_sersic
        )

        amplitudes, sigmas, norm = mge.mge_1d(rs, ss, N=n_comp)
        ss_mge = self.multiGaussian.function(
            rs, np.zeros_like(rs), amp=amplitudes, sigma=sigmas
        )
        # print((ss - ss_mge)/ss)
        for i in range(10, len(ss) - 10):
            # print(rs[i])
            npt.assert_almost_equal((ss_mge[i] - ss[i]) / ss[i], 0, decimal=1)

        amplitudes, sigmas, norm = mge.mge_1d(rs, np.zeros_like(rs), N=n_comp)
        assert amplitudes[0] == 0

        amplitudes, sigmas, norm = mge.mge_1d(rs, np.zeros_like(rs), N=0)
        assert amplitudes[0] == 0

    def test_mge_sersic_radius(self):
        n_comp = 30
        r_sersic = 0.5
        n_sersic = 3.7
        I0_sersic = 1.0
        rs = np.logspace(-2.0, 1.0, 50) * r_sersic
        ss = self.sersic.function(
            rs, np.zeros_like(rs), amp=I0_sersic, n_sersic=n_sersic, R_sersic=r_sersic
        )

        amplitudes, sigmas, norm = mge.mge_1d(rs, ss, N=n_comp)
        ss_mge = self.multiGaussian.function(
            rs, np.zeros_like(rs), amp=amplitudes, sigma=sigmas
        )
        print((ss - ss_mge) / (ss + ss_mge))
        for i in range(10, len(ss) - 10):
            # print(rs[i])
            npt.assert_almost_equal((ss_mge[i] - ss[i]) / (ss[i]), 0, decimal=1)

    def test_mge_sersic_n_sersic(self):
        n_comp = 20
        r_sersic = 1.5
        n_sersic = 0.5
        I0_sersic = 1.0
        rs = np.logspace(-2.0, 1.0, 50) * r_sersic
        ss = self.sersic.function(
            rs, np.zeros_like(rs), amp=I0_sersic, n_sersic=n_sersic, R_sersic=r_sersic
        )

        amplitudes, sigmas, norm = mge.mge_1d(rs, ss, N=n_comp)
        ss_mge = self.multiGaussian.function(
            rs, np.zeros_like(rs), amp=amplitudes, sigma=sigmas
        )
        for i in range(10, len(ss) - 10):
            npt.assert_almost_equal(
                (ss_mge[i] - ss[i]) / (ss[i] + ss_mge[i]), 0, decimal=1
            )

        n_comp = 20
        r_sersic = 1.5
        n_sersic = 3.5
        I0_sersic = 1.0
        rs = np.logspace(-2.0, 1.0, 50) * r_sersic
        ss = self.sersic.function(
            rs, np.zeros_like(rs), amp=I0_sersic, n_sersic=n_sersic, R_sersic=r_sersic
        )

        amplitudes, sigmas, norm = mge.mge_1d(rs, ss, N=n_comp)
        ss_mge = self.multiGaussian.function(
            rs, np.zeros_like(rs), amp=amplitudes, sigma=sigmas
        )
        for i in range(10, len(ss) - 10):
            npt.assert_almost_equal(
                (ss_mge[i] - ss[i]) / (ss[i] + ss_mge[i]), 0, decimal=1
            )

    def test_hernquist(self):
        hernquist = Hernquist()
        n_comp = 20
        sigma0 = 1
        r_eff = 1.5
        rs = np.logspace(-2.0, 1.0, 50) * r_eff * 0.5
        ss = hernquist.function(rs, np.zeros_like(rs), sigma0, Rs=r_eff)
        amplitudes, sigmas, norm = mge.mge_1d(rs, ss, N=n_comp)
        ss_mge = self.multiGaussian.function(
            rs, np.zeros_like(rs), amp=amplitudes, sigma=sigmas
        )
        for i in range(10, len(ss) - 10):
            npt.assert_almost_equal(
                (ss_mge[i] - ss[i]) / (ss[i] + ss_mge[i]), 0, decimal=2
            )

    def test_hernquist_deprojection(self):
        hernquist = Hernquist()
        n_comp = 20
        sigma0 = 1
        r_eff = 1.5
        rs = np.logspace(-2.0, 1.0, 50) * r_eff * 0.5
        ss = hernquist.function(rs, np.zeros_like(rs), sigma0, Rs=r_eff)
        amplitudes, sigmas, norm = mge.mge_1d(rs, ss, N=n_comp)
        amplitudes_3d, sigmas_3d = mge.de_projection_3d(amplitudes, sigmas)
        ss_3d_mge = self.multiGaussian.function(
            rs, np.zeros_like(rs), amp=amplitudes_3d, sigma=sigmas_3d
        )
        ss_3d_mulit = self.multiGaussian.light_3d(rs, amp=amplitudes, sigma=sigmas)
        for i in range(10, len(ss_3d_mge)):
            npt.assert_almost_equal(
                (ss_3d_mge[i] - ss_3d_mulit[i]) / (ss_3d_mulit[i] + ss_3d_mge[i]),
                0,
                decimal=1,
            )

        ss_3d = hernquist.light_3d(rs, sigma0, Rs=r_eff)
        for i in range(10, len(ss_3d) - 10):
            npt.assert_almost_equal(
                (ss_3d_mge[i] - ss_3d[i]) / (ss_3d[i] + ss_3d_mge[i]), 0, decimal=1
            )

    def test_spemd(self):
        from lenstronomy.LensModel.Profiles.spep import SPEP
        from lenstronomy.LensModel.Profiles.multi_gaussian_kappa import (
            MultiGaussianKappa,
        )

        spep = SPEP()
        mge_kappa = MultiGaussianKappa()
        n_comp = 8
        theta_E = 1.41
        kwargs = {"theta_E": theta_E, "e1": 0, "e2": 0, "gamma": 1.61}
        rs = np.logspace(-2.0, 1.0, 100) * theta_E
        f_xx, f_xy, f_yx, f_yy = spep.hessian(rs, 0, **kwargs)
        kappa = 1 / 2.0 * (f_xx + f_yy)
        amplitudes, sigmas, norm = mge.mge_1d(rs, kappa, N=n_comp)
        kappa_mge = self.multiGaussian.function(
            rs, np.zeros_like(rs), amp=amplitudes, sigma=sigmas
        )
        f_xx_mge, f_xy_mge, f_yx_mge, f_yy_mge = mge_kappa.hessian(
            rs, np.zeros_like(rs), amp=amplitudes, sigma=sigmas
        )
        for i in range(0, 80):
            npt.assert_almost_equal(
                kappa_mge[i], 1.0 / 2 * (f_xx_mge[i] + f_yy_mge[i]), decimal=1
            )
            npt.assert_almost_equal((kappa[i] - kappa_mge[i]) / kappa[i], 0, decimal=1)

        f_ = spep.function(theta_E, 0, **kwargs)
        f_mge = mge_kappa.function(theta_E, 0, sigma=sigmas, amp=amplitudes)
        npt.assert_almost_equal(f_mge / f_, 1, decimal=2)

    def test_example(self):
        n_comp = 10
        rs = np.array(
            [
                0.01589126,
                0.01703967,
                0.01827108,
                0.01959148,
                0.0210073,
                0.02252544,
                0.02415329,
                0.02589879,
                0.02777042,
                0.02977731,
                0.03192923,
                0.03423667,
                0.03671086,
                0.03936385,
                0.04220857,
                0.04525886,
                0.0485296,
                0.0520367,
                0.05579724,
                0.05982956,
                0.06415327,
                0.06878945,
                0.07376067,
                0.07909115,
                0.08480685,
                0.09093561,
                0.09750727,
                0.10455385,
                0.11210966,
                0.12021152,
                0.12889887,
                0.13821403,
                0.14820238,
                0.15891255,
                0.17039672,
                0.18271082,
                0.19591482,
                0.21007304,
                0.22525444,
                0.24153295,
                0.25898787,
                0.2777042,
                0.29777311,
                0.31929235,
                0.34236672,
                0.36710861,
                0.39363853,
                0.42208569,
                0.45258865,
                0.48529597,
                0.52036697,
                0.55797244,
                0.59829556,
                0.64153272,
                0.6878945,
                0.73760673,
                0.79091152,
                0.8480685,
                0.90935605,
                0.97507269,
                1.04553848,
                1.12109664,
                1.20211518,
                1.28898871,
                1.38214034,
                1.48202378,
                1.58912553,
                1.70396721,
                1.82710819,
                1.95914822,
                2.10073042,
                2.25254437,
                2.4153295,
                2.58987865,
                2.77704199,
                2.9777311,
                3.19292345,
                3.42366716,
                3.67108607,
                3.93638527,
                4.22085689,
                4.5258865,
                4.85295974,
                5.20366966,
                5.57972441,
                5.98295559,
                6.41532717,
                6.87894505,
                7.37606729,
                7.90911519,
                8.48068497,
                9.09356051,
                9.75072687,
                10.45538481,
                11.21096643,
                12.02115183,
                12.88988708,
                13.82140341,
                14.82023784,
                15.89125526,
            ]
        )
        kappa = np.array(
            [
                12.13776067,
                11.60484966,
                11.09533396,
                10.60818686,
                10.14242668,
                9.69711473,
                9.27135349,
                8.86428482,
                8.47508818,
                8.10297905,
                7.7472073,
                7.40705574,
                7.08183863,
                6.77090034,
                6.47361399,
                6.18938022,
                5.917626,
                5.65780342,
                5.40938864,
                5.1718808,
                4.94480104,
                4.72769151,
                4.52011448,
                4.3216514,
                4.13190214,
                3.9504841,
                3.77703149,
                3.61119459,
                3.45263901,
                3.30104507,
                3.1561071,
                3.01753287,
                2.88504297,
                2.75837025,
                2.63725931,
                2.52146595,
                2.41075668,
                2.30490829,
                2.20370736,
                2.10694982,
                2.01444058,
                1.92599312,
                1.84142909,
                1.76057799,
                1.6832768,
                1.60936965,
                1.53870751,
                1.47114792,
                1.40655465,
                1.34479745,
                1.28575181,
                1.22929867,
                1.17532421,
                1.12371958,
                1.07438074,
                1.02720821,
                0.98210687,
                0.93898578,
                0.897758,
                0.85834039,
                0.82065349,
                0.78462129,
                0.75017114,
                0.71723359,
                0.68574222,
                0.65563353,
                0.62684681,
                0.59932403,
                0.57300967,
                0.5478507,
                0.52379638,
                0.5007982,
                0.47880979,
                0.45778683,
                0.43768691,
                0.41846951,
                0.40009589,
                0.38252899,
                0.3657334,
                0.34967525,
                0.33432216,
                0.31964317,
                0.30560868,
                0.29219041,
                0.27936129,
                0.26709545,
                0.25536817,
                0.24415579,
                0.23343571,
                0.22318631,
                0.21338694,
                0.20401782,
                0.19506006,
                0.18649562,
                0.17830721,
                0.17047832,
                0.16299318,
                0.15583668,
                0.14899441,
                0.14245255,
            ]
        )
        amplitudes, sigmas, norm = mge.mge_1d(rs, kappa, N=n_comp)

    def test_nfw_sersic(self):
        kwargs_lens_nfw = {"alpha_Rs": 1.4129647849966354, "Rs": 7.0991113634274736}
        kwargs_lens_sersic = {
            "k_eff": 0.24100561407593576,
            "n_sersic": 1.8058507329346063,
            "R_sersic": 1.0371803141813705,
        }
        from lenstronomy.LensModel.Profiles.nfw import NFW
        from lenstronomy.LensModel.Profiles.sersic import Sersic

        nfw = NFW()
        sersic = Sersic()
        theta_E = 1.5
        n_comp = 10
        rs = np.logspace(-2.0, 1.0, 100) * theta_E
        f_xx_nfw, f_xy_nfw, f_yx_nfw, f_yy_nfw = nfw.hessian(rs, 0, **kwargs_lens_nfw)
        f_xx_s, f_xy_s, f_yx_s, f_yy_s = sersic.hessian(rs, 0, **kwargs_lens_sersic)
        kappa = 1 / 2.0 * (f_xx_nfw + f_xx_s + f_yy_nfw + f_yy_s)
        amplitudes, sigmas, norm = mge.mge_1d(rs, kappa, N=n_comp)
        kappa_mge = self.multiGaussian.function(
            rs, np.zeros_like(rs), amp=amplitudes, sigma=sigmas
        )
        from lenstronomy.LensModel.Profiles.multi_gaussian_kappa import (
            MultiGaussianKappa,
        )

        mge_kappa = MultiGaussianKappa()
        f_xx_mge, f_xy_mge, f_yx_mge, f_yy_mge = mge_kappa.hessian(
            rs, np.zeros_like(rs), amp=amplitudes, sigma=sigmas
        )
        for i in range(0, 80):
            npt.assert_almost_equal(
                kappa_mge[i], 1.0 / 2 * (f_xx_mge[i] + f_yy_mge[i]), decimal=1
            )
            npt.assert_almost_equal((kappa[i] - kappa_mge[i]) / kappa[i], 0, decimal=1)

        f_nfw = nfw.function(theta_E, 0, **kwargs_lens_nfw)
        f_s = sersic.function(theta_E, 0, **kwargs_lens_sersic)
        f_mge = mge_kappa.function(theta_E, 0, sigma=sigmas, amp=amplitudes)
        npt.assert_almost_equal(f_mge / (f_nfw + f_s), 1, decimal=2)


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