__author__ = "sibirrer"

import numpy.testing as npt
import numpy as np
import pytest
from lenstronomy.LensModel.Profiles.curved_arc_tan_diff import CurvedArcTanDiff
from lenstronomy.LensModel.Profiles.sie import SIE
from lenstronomy.LensModel.Profiles.convergence import Convergence
from lenstronomy.LensModel.lens_model import LensModel
from lenstronomy.LensModel.lens_model_extensions import LensModelExtensions


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

    def setup_method(self):
        self.model = CurvedArcTanDiff()
        self.sie = SIE()
        self.mst = Convergence()

    def test_curved_arc_round_recovery(self):
        """Test whether the curved arc parameters are satisfied in differential form."""

        center_x, center_y = (
            2,
            0.0,
        )  # test works except at (0,0) where the direction angle is not well defined
        tangential_stretch = 4.0
        radial_stretch = 1.0
        curvature, direction = 0.5, 0.5
        dtan_dtan = 0
        kwargs_lens = {
            "tangential_stretch": tangential_stretch,
            "radial_stretch": radial_stretch,
            "curvature": curvature,
            "direction": direction,
            "dtan_dtan": dtan_dtan,
            "center_x": center_x,
            "center_y": center_y,
        }
        self._test_curved_arc_recovery(kwargs_lens)

        kwargs_lens = {
            "tangential_stretch": tangential_stretch,
            "radial_stretch": 1.1,
            "curvature": curvature,
            "direction": 0.01,
            "dtan_dtan": dtan_dtan,
            "center_x": center_x,
            "center_y": center_y,
        }
        self._test_curved_arc_recovery(kwargs_lens)

        kwargs_lens = {
            "tangential_stretch": 10,
            "radial_stretch": 1.0,
            "curvature": 0.2,
            "direction": 0.01,
            "dtan_dtan": 0.0,
            "center_x": center_x,
            "center_y": center_y,
        }
        self._test_curved_arc_recovery(kwargs_lens)

    def test_curved_arc_recovery(self):
        """Test whether the curved arc parameters are satisfied in differential form."""

        center_x, center_y = (
            3,
            0,
        )  # test works except at (0,0) where the direction angle is not well defined

        kwargs_lens = {
            "tangential_stretch": 2.0,
            "radial_stretch": 1.0,
            "curvature": 0.3,
            "direction": 0.001,
            "dtan_dtan": 0.1,
            "center_x": center_x,
            "center_y": center_y,
        }
        self._test_curved_arc_recovery(kwargs_lens)

        # and here we change directions
        kwargs_lens = {
            "tangential_stretch": 2.0,
            "radial_stretch": 1.0,
            "curvature": 0.3,
            "direction": 0.5,
            "dtan_dtan": 0.1,
            "center_x": center_x,
            "center_y": center_y,
        }
        self._test_curved_arc_recovery(kwargs_lens)

        # and here we have the radial stretch != 1, thus applying an MST
        kwargs_lens = {
            "tangential_stretch": 2.0,
            "radial_stretch": 1.1,
            "curvature": 0.3,
            "direction": 0.5,
            "dtan_dtan": 0.1,
            "center_x": center_x,
            "center_y": center_y,
        }
        self._test_curved_arc_recovery(kwargs_lens)

        kwargs_lens = {
            "tangential_stretch": 2.0,
            "radial_stretch": 1.1,
            "curvature": 0.3,
            "direction": 0.5,
            "dtan_dtan": -0.1,
            "center_x": center_x,
            "center_y": center_y,
        }
        self._test_curved_arc_recovery(kwargs_lens)

    def _test_in_out_scaling(self):
        # some scaling tests with plots that are going to be ignored
        ext = LensModelExtensions(LensModel(lens_model_list=["CURVED_ARC_TAN_DIFF"]))

        # change in dtan_dtan for fixed other components input vs output comparison
        dtan_dtan_list = np.linspace(start=-0.5, stop=0.5, num=21)
        dtan_dtan_out_list = []
        kwargs_lens = {
            "tangential_stretch": 2.0,
            "radial_stretch": 1.0,
            "curvature": 0.2,
            "direction": 0.001,
            "dtan_dtan": 0.1,
            "center_x": 1,
            "center_y": 0,
        }
        for dtan_dtan in dtan_dtan_list:
            kwargs_lens["dtan_dtan"] = dtan_dtan
            (
                lambda_rad,
                lambda_tan,
                orientation_angle,
                dlambda_tan_dtan,
                dlambda_tan_drad,
                dlambda_rad_drad,
                dlambda_rad_dtan,
                dphi_tan_dtan,
                dphi_tan_drad,
                dphi_rad_drad,
                dphi_rad_dtan,
            ) = ext.radial_tangential_differentials(1, 0, [kwargs_lens])
            dtan_dtan_out_list.append(dlambda_tan_dtan)
        dtan_dtan_out_list = np.array(dtan_dtan_out_list)
        import matplotlib.pyplot as plt

        plt.plot(dtan_dtan_list, dtan_dtan_out_list)
        plt.xlabel("dtan_in fixed lens")
        plt.ylabel("dtan_out")
        plt.show()
        # npt.assert_almost_equal(dtan_dtan_out_list, dtan_dtan_list)

        # change in tangential stretch
        lambda_tan_list = np.linspace(start=2, stop=20, num=21)
        dtan_dtan_out_list = []
        kwargs_lens = {
            "tangential_stretch": 2.0,
            "radial_stretch": 1.0,
            "curvature": 0.2,
            "direction": 0.001,
            "dtan_dtan": 0.1,
            "center_x": 1,
            "center_y": 0,
        }
        for lambda_tan in lambda_tan_list:
            kwargs_lens["tangential_stretch"] = lambda_tan
            (
                lambda_rad,
                lambda_tan,
                orientation_angle,
                dlambda_tan_dtan,
                dlambda_tan_drad,
                dlambda_rad_drad,
                dlambda_rad_dtan,
                dphi_tan_dtan,
                dphi_tan_drad,
                dphi_rad_drad,
                dphi_rad_dtan,
            ) = ext.radial_tangential_differentials(1, 0, [kwargs_lens])
            dtan_dtan_out_list.append(dlambda_tan_dtan)
        dtan_dtan_out_list = np.array(dtan_dtan_out_list)
        import matplotlib.pyplot as plt

        plt.plot(lambda_tan_list, dtan_dtan_out_list)
        plt.xlabel("lambda_tan")
        plt.ylabel("dtan_out")
        plt.show()
        # npt.assert_almost_equal(dtan_dtan_out_list, dtan_dtan_list)

        # change in curvature radius
        curvature_list = np.linspace(start=0.1, stop=1, num=21)
        dtan_dtan_out_list = []
        kwargs_lens = {
            "tangential_stretch": 5.0,
            "radial_stretch": 1.0,
            "curvature": 0.2,
            "direction": 0.001,
            "dtan_dtan": 0.1,
            "center_x": 1,
            "center_y": 0,
        }
        for curvatrue in curvature_list:
            kwargs_lens["curvature"] = curvatrue
            (
                lambda_rad,
                lambda_tan,
                orientation_angle,
                dlambda_tan_dtan,
                dlambda_tan_drad,
                dlambda_rad_drad,
                dlambda_rad_dtan,
                dphi_tan_dtan,
                dphi_tan_drad,
                dphi_rad_drad,
                dphi_rad_dtan,
            ) = ext.radial_tangential_differentials(1, 0, [kwargs_lens])
            dtan_dtan_out_list.append(dlambda_tan_dtan)
        dtan_dtan_out_list = np.array(dtan_dtan_out_list)
        import matplotlib.pyplot as plt

        plt.plot(lambda_tan_list, dtan_dtan_out_list)
        plt.xlabel("curvature")
        plt.ylabel("dtan_out")
        plt.show()
        # npt.assert_almost_equal(dtan_dtan_out_list, dtan_dtan_list)

    def _test_curved_arc_recovery(self, kwargs_arc_init):
        ext = LensModelExtensions(LensModel(lens_model_list=["CURVED_ARC_TAN_DIFF"]))
        center_x, center_y = kwargs_arc_init["center_x"], kwargs_arc_init["center_y"]
        kwargs_arc = ext.curved_arc_estimate(center_x, center_y, [kwargs_arc_init])
        (
            lambda_rad,
            lambda_tan,
            orientation_angle,
            dlambda_tan_dtan,
            dlambda_tan_drad,
            dlambda_rad_drad,
            dlambda_rad_dtan,
            dphi_tan_dtan,
            dphi_tan_drad,
            dphi_rad_drad,
            dphi_rad_dtan,
        ) = ext.radial_tangential_differentials(center_x, center_y, [kwargs_arc_init])
        print(lambda_tan, dlambda_tan_dtan, kwargs_arc_init["dtan_dtan"])
        npt.assert_almost_equal(
            kwargs_arc["tangential_stretch"] / kwargs_arc_init["tangential_stretch"],
            1,
            decimal=3,
        )
        npt.assert_almost_equal(
            kwargs_arc["radial_stretch"], kwargs_arc_init["radial_stretch"], decimal=3
        )
        npt.assert_almost_equal(
            kwargs_arc["curvature"], kwargs_arc_init["curvature"], decimal=2
        )
        npt.assert_almost_equal(dphi_tan_dtan, kwargs_arc_init["curvature"], decimal=2)
        npt.assert_almost_equal(
            kwargs_arc["direction"], kwargs_arc_init["direction"], decimal=3
        )
        npt.assert_almost_equal(
            dlambda_tan_dtan / lambda_tan, kwargs_arc_init["dtan_dtan"], decimal=2
        )


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