__author__ = "sibirrer"

import numpy as np
import numpy.testing as npt
import pytest
from lenstronomy.LensModel.Profiles.curved_arc_sis_mst import CurvedArcSISMST
from lenstronomy.LensModel.Profiles.curved_arc_const import (
    CurvedArcConstMST,
    CurvedArcConst,
)
from lenstronomy.Util import util


class TestCurvedArcConst(object):
    """Tests the source model routines."""

    def setup_method(self):
        self.arc_sis = CurvedArcSISMST()
        self.arc_const = CurvedArcConst()

    def test_function(self):
        kwargs_arc = {
            "tangential_stretch": 5,
            #'radial_stretch': 1,
            "curvature": 1.0 / 10,
            "direction": 0,
            "center_x": 0,
            "center_y": 0,
        }
        npt.assert_raises(Exception, self.arc_const.function, 0.0, 0.0, **kwargs_arc)

    def test_derivatives(self):
        kwargs_arc = {
            "tangential_stretch": 3,
            #'radial_stretch': 1.,
            "curvature": 0.8,
            "direction": 0,
            "center_x": 0,
            "center_y": 0,
        }

        kwargs_arc_sis = {
            "tangential_stretch": 3,
            "radial_stretch": 1.0,
            "curvature": 0.8,
            "direction": 0,
            "center_x": 0,
            "center_y": 0,
        }
        x, y = util.make_grid(numPix=100, deltapix=0.01)
        f_x_sis, f_y_sis = self.arc_sis.derivatives(x, y, **kwargs_arc_sis)
        beta_x_sis = x - f_x_sis
        beta_y_sis = y - f_y_sis
        f_x_const, f_y_const = self.arc_const.derivatives(x, y, **kwargs_arc)
        beta_x_const = x - f_x_const
        beta_y_const = y - f_y_const

        from lenstronomy.LightModel.light_model import LightModel

        gauss = LightModel(["GAUSSIAN"])
        kwargs_source = [{"amp": 1, "sigma": 0.05, "center_x": 0, "center_y": 0}]
        flux_sis = gauss.surface_brightness(beta_x_sis, beta_y_sis, kwargs_source)
        flux_const = gauss.surface_brightness(beta_x_const, beta_y_const, kwargs_source)

        npt.assert_almost_equal(
            (flux_const - flux_sis) / np.max(flux_const), 0, decimal=2
        )

        # check for stability outside the defined bounds of curvature
        f_x_const, f_y_const = self.arc_const.derivatives(x=0, y=1000, **kwargs_arc)
        npt.assert_almost_equal(f_x_const, 0)
        npt.assert_almost_equal(f_y_const, 0)

    def test_hessian(self):
        kwargs_arc = {
            "tangential_stretch": 5,
            #'radial_stretch': 1,
            "curvature": 1.0 / 10,
            "direction": 0.5,
            "center_x": 0,
            "center_y": 0,
        }
        x, y = 0.0, 1.0
        f_xx, f_xy, f_yx, f_yy = self.arc_const.hessian(x, y, **kwargs_arc)

        alpha_ra, alpha_dec = self.arc_const.derivatives(x, y, **kwargs_arc)
        diff = 0.0000001
        alpha_ra_dx, alpha_dec_dx = self.arc_const.derivatives(
            x + diff, y, **kwargs_arc
        )
        alpha_ra_dy, alpha_dec_dy = self.arc_const.derivatives(
            x, y + diff, **kwargs_arc
        )

        f_xx_num = (alpha_ra_dx - alpha_ra) / diff
        f_xy_num = (alpha_ra_dy - alpha_ra) / diff
        f_yx_num = (alpha_dec_dx - alpha_dec) / diff
        f_yy_num = (alpha_dec_dy - alpha_dec) / diff
        print(f_xx, f_xx_num)
        print(f_xy, f_xy_num)
        print(f_yx, f_yx_num)
        print(f_yy, f_yy_num)

        npt.assert_almost_equal(f_xx, f_xx_num)
        npt.assert_almost_equal(f_xy, f_xy_num)
        npt.assert_almost_equal(f_yx, f_yx_num)
        npt.assert_almost_equal(f_yy, f_yy_num)


class TestCurvedArcConstMST(object):
    def setup_method(self):
        self.arc_const = CurvedArcConstMST()

    def test_function(self):
        kwargs_arc = {
            "tangential_stretch": 5,
            "radial_stretch": 1,
            "curvature": 1.0 / 10,
            "direction": 0,
            "center_x": 0,
            "center_y": 0,
        }
        npt.assert_raises(Exception, self.arc_const.function, 0.0, 0.0, **kwargs_arc)

    def test_hessian(self):
        kwargs_arc = {
            "tangential_stretch": 5,
            "radial_stretch": 1,
            "curvature": 1.0 / 10,
            "direction": 0,
            "center_x": 0,
            "center_y": 0,
        }
        # npt.assert_raises(Exception, self.arc_const.hessian, 0., 0., **kwargs_arc)


if __name__ == "__main__":
    pytest.main("-k TestLensModel")
