__author__ = "sibirrer"


from lenstronomy.LightModel.Profiles.gaussian import Gaussian
from lenstronomy.LensModel.Profiles.gaussian_kappa import GaussianKappa
from lenstronomy.LensModel.Profiles.gaussian_ellipse_potential import (
    GaussianEllipsePotential,
)
from lenstronomy.LensModel.Profiles.multi_gaussian_kappa import (
    MultiGaussianKappa,
    MultiGaussianKappaEllipse,
)

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


class TestGaussianKappa(object):
    """Test the Gaussian with Gaussian kappa."""

    def setup_method(self):
        self.gaussian_kappa = MultiGaussianKappa()
        self.gaussian = Gaussian()
        self.g_kappa = GaussianKappa()

    def test_derivatives(self):
        x = np.linspace(0, 5, 10)
        y = np.linspace(0, 5, 10)
        amp = [1.0 * 2 * np.pi]
        center_x = 0.0
        center_y = 0.0
        sigma = [1.0]
        f_x, f_y = self.gaussian_kappa.derivatives(x, y, amp, sigma, center_x, center_y)
        npt.assert_almost_equal(f_x[2], 0.63813558702212059, decimal=8)
        npt.assert_almost_equal(f_y[2], 0.63813558702212059, decimal=8)

    def test_hessian(self):
        x = np.linspace(0, 5, 10)
        y = np.linspace(0, 5, 10)
        amp = [1.0 * 2 * np.pi]
        center_x = 0.0
        center_y = 0.0
        sigma = [1.0]
        f_xx, f_xy, f_yx, f_yy = self.gaussian_kappa.hessian(
            x, y, amp, sigma, center_x, center_y
        )
        kappa = 1.0 / 2 * (f_xx + f_yy)
        kappa_true = self.gaussian.function(x, y, amp[0], sigma[0], center_x, center_y)
        print(kappa_true)
        print(kappa)
        npt.assert_almost_equal(kappa[0], kappa_true[0], decimal=5)
        npt.assert_almost_equal(kappa[1], kappa_true[1], decimal=5)
        npt.assert_almost_equal(f_xy, f_yx, decimal=8)

    def test_density_2d(self):
        x = np.linspace(0, 5, 10)
        y = np.linspace(0, 5, 10)
        amp = [1.0 * 2 * np.pi]
        center_x = 0.0
        center_y = 0.0
        sigma = [1.0]
        f_xx, f_xy, f_yx, f_yy = self.gaussian_kappa.hessian(
            x, y, amp, sigma, center_x, center_y
        )
        kappa = 1.0 / 2 * (f_xx + f_yy)
        amp_3d = self.g_kappa._amp2d_to_3d(amp, sigma[0], sigma[0])
        density_2d = self.gaussian_kappa.density_2d(
            x, y, amp_3d, sigma, center_x, center_y
        )
        npt.assert_almost_equal(kappa[1], density_2d[1], decimal=5)
        npt.assert_almost_equal(kappa[2], density_2d[2], decimal=5)

    def test_density(self):
        amp = [1.0 * 2 * np.pi]

        sigma = [1.0]
        density = self.gaussian_kappa.density(1.0, amp, sigma)
        npt.assert_almost_equal(density, 0.6065306597126334, decimal=8)


class TestGaussianKappaEllipse(object):
    """Test the Gaussian with Gaussian kappa."""

    def setup_method(self):
        self.multi = MultiGaussianKappaEllipse()
        self.single = GaussianEllipsePotential()

    def test_function(self):
        x, y = 1, 2
        amp = 1
        sigma = 1
        e1, e2 = 0.1, -0.1
        center_x, center_y = 1, 0
        f_ = self.multi.function(
            x,
            y,
            amp=[amp],
            sigma=[sigma],
            e1=e1,
            e2=e2,
            center_x=center_x,
            center_y=center_y,
        )
        f_single = self.single.function(
            x,
            y,
            amp=amp,
            sigma=sigma,
            e1=e1,
            e2=e2,
            center_x=center_x,
            center_y=center_y,
        )
        npt.assert_almost_equal(f_, f_single, decimal=8)

    def test_derivatives(self):
        x, y = 1, 2
        amp = 1
        sigma = 1
        e1, e2 = 0.1, -0.1
        center_x, center_y = 1, 0
        f_x, f_y = self.multi.derivatives(
            x,
            y,
            amp=[amp],
            sigma=[sigma],
            e1=e1,
            e2=e2,
            center_x=center_x,
            center_y=center_y,
        )
        f_x_s, f_y_s = self.single.derivatives(
            x,
            y,
            amp=amp,
            sigma=sigma,
            e1=e1,
            e2=e2,
            center_x=center_x,
            center_y=center_y,
        )
        npt.assert_almost_equal(f_x, f_x_s, decimal=8)
        npt.assert_almost_equal(f_y, f_y_s, decimal=8)

    def test_hessian(self):
        x, y = 1, 2
        amp = 1
        sigma = 1
        e1, e2 = 0.1, -0.1
        center_x, center_y = 1, 0
        f_xx, f_xy, f_yx, f_yy = self.multi.hessian(
            x,
            y,
            amp=[amp],
            sigma=[sigma],
            e1=e1,
            e2=e2,
            center_x=center_x,
            center_y=center_y,
        )
        f_xx_s, f_xy_s, f_yx_s, f_yy_s = self.single.hessian(
            x,
            y,
            amp=amp,
            sigma=sigma,
            e1=e1,
            e2=e2,
            center_x=center_x,
            center_y=center_y,
        )
        npt.assert_almost_equal(f_xx, f_xx_s, decimal=8)
        npt.assert_almost_equal(f_yy, f_yy_s, decimal=8)
        npt.assert_almost_equal(f_xy, f_xy_s, decimal=8)
        npt.assert_almost_equal(f_yx, f_yx_s, decimal=7)

    def test_density_2d(self):
        x, y = 1, 2
        amp = 1
        sigma = 1
        e1, e2 = 0.1, -0.1
        center_x, center_y = 1, 0
        f_ = self.multi.density_2d(
            x,
            y,
            amp=[amp],
            sigma=[sigma],
            e1=e1,
            e2=e2,
            center_x=center_x,
            center_y=center_y,
        )
        f_single = self.single.density_2d(
            x,
            y,
            amp=amp,
            sigma=sigma,
            e1=e1,
            e2=e2,
            center_x=center_x,
            center_y=center_y,
        )
        npt.assert_almost_equal(f_, f_single, decimal=8)

    def test_density(self):
        r = 1
        amp = 1
        sigma = 1
        e1, e2 = 0.1, -0.1
        f_ = self.multi.density(r, amp=[amp], sigma=[sigma], e1=e1, e2=e2)
        f_single = self.single.density(r, amp=amp, sigma=sigma, e1=e1, e2=e2)
        npt.assert_almost_equal(f_, f_single, decimal=8)


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