__author__ = "sibirrer"


from lenstronomy.LensModel.Profiles.cored_density_mst import CoredDensityMST

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


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

    def setup_method(self):
        self.profile1 = CoredDensityMST(profile_type="CORED_DENSITY")
        self.profile2 = CoredDensityMST(profile_type="CORED_DENSITY_2")
        self.profile3 = CoredDensityMST(profile_type="CORED_DENSITY_EXP")
        self.profile4 = CoredDensityMST(profile_type="CORED_DENSITY_ULDM")
        self.kwargs_lens = {
            "lambda_approx": 0.9,
            "r_core": 100,
            "center_x": 0,
            "center_y": 0,
        }

    def test_function(self):
        x = np.array([0.01, 1])
        y = np.array([0, 0])
        f_ = self.profile1.function(x, y, **self.kwargs_lens)
        npt.assert_almost_equal(
            f_[0] - f_[1], 0, decimal=3
        )  # test to demand that the profile is (almost) zero
        f_ = self.profile2.function(x, y, **self.kwargs_lens)
        npt.assert_almost_equal(
            f_[0] - f_[1], 0, decimal=3
        )  # test to demand that the profile is (almost) zero
        f_ = self.profile3.function(x, y, **self.kwargs_lens)
        npt.assert_almost_equal(
            f_[0] - f_[1], 0, decimal=3
        )  # test to demand that the profile is (almost) zero
        f_ = self.profile4.function(x, y, **self.kwargs_lens)
        npt.assert_almost_equal(
            f_[0] - f_[1], 0, decimal=3
        )  # test to demand that the profile is (almost) zero

    def test_derivatives(self):
        x = np.array([0.01, 1])
        y = np.array([0, 0])
        f_x, f_y = self.profile1.derivatives(x, y, **self.kwargs_lens)
        npt.assert_almost_equal(
            f_x[0] - f_x[1], 0, decimal=3
        )  # test to demand that the profile is (almost) zero
        f_x, f_y = self.profile2.derivatives(x, y, **self.kwargs_lens)
        npt.assert_almost_equal(
            f_x[0] - f_x[1], 0, decimal=3
        )  # test to demand that the profile is (almost) zero
        f_x, f_y = self.profile3.derivatives(x, y, **self.kwargs_lens)
        npt.assert_almost_equal(
            f_x[0] - f_x[1], 0, decimal=3
        )  # test to demand that the profile is (almost) zero
        f_x, f_y = self.profile4.derivatives(x, y, **self.kwargs_lens)
        npt.assert_almost_equal(
            f_x[0] - f_x[1], 0, decimal=3
        )  # test to demand that the profile is (almost) zero

    def test_hessian(self):
        x = np.array([0.01, 1])
        y = np.array([0, 0])
        f_xx, f_xy, f_yx, f_yy = self.profile1.hessian(x, y, **self.kwargs_lens)
        npt.assert_almost_equal(
            f_xx, 0, decimal=3
        )  # test to demand that the profile is (almost) zero
        f_xx, f_xy, f_yx, f_yy = self.profile2.hessian(x, y, **self.kwargs_lens)
        npt.assert_almost_equal(
            f_xx, 0, decimal=3
        )  # test to demand that the profile is (almost) zero
        f_xx, f_xy, f_yx, f_yy = self.profile3.hessian(x, y, **self.kwargs_lens)
        npt.assert_almost_equal(
            f_xx, 0, decimal=3
        )  # test to demand that the profile is (almost) zero
        f_xx, f_xy, f_yx, f_yy = self.profile4.hessian(x, y, **self.kwargs_lens)
        npt.assert_almost_equal(
            f_xx, 0, decimal=3
        )  # test to demand that the profile is (almost) zero


class TestRaise(unittest.TestCase):
    def test_raise(self):
        with self.assertRaises(ValueError):
            CoredDensityMST(profile_type="WRONG_PROFILE")


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