__author__ = "sibirrer"

import numpy as np
import numpy.testing as npt
import pytest
import unittest
from lenstronomy.LensModel.Profiles.nfw import NFW
from lenstronomy.LensModel.lens_model import LensModel
from lenstronomy.LensModel.Solver.lens_equation_solver import LensEquationSolver
from lenstronomy.LensModel.Solver.solver2point import Solver2Point
import lenstronomy.Util.param_util as param_util


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

    def setup_method(self):
        pass

    def test_subtract(self):
        lensModel = LensModel(["SPEP"])
        solver_spep_center = Solver2Point(lensModel, solver_type="CENTER")
        x_cat = np.array([0, 0])
        y_cat = np.array([1, 2])
        a = solver_spep_center._subtract_constraint(x_cat, y_cat)
        assert a[0] == 0
        assert a[1] == 1

    def test_all_spep(self):
        lensModel = LensModel(["SPEP"])
        solver_spep_center = Solver2Point(lensModel, solver_type="CENTER")
        solver_spep_ellipse = Solver2Point(lensModel, solver_type="ELLIPSE")
        image_position_spep = LensEquationSolver(lensModel)
        sourcePos_x = 0.1
        sourcePos_y = 0.03
        gamma = 1.9
        phi_G, q = 0.5, 0.8
        e1, e2 = param_util.phi_q2_ellipticity(phi_G, q)
        kwargs_lens = [
            {
                "theta_E": 1,
                "gamma": gamma,
                "e1": e1,
                "e2": e2,
                "center_x": 0.1,
                "center_y": -0.1,
            }
        ]
        x_pos, y_pos = image_position_spep.findBrightImage(
            sourcePos_x,
            sourcePos_y,
            kwargs_lens,
            numImages=2,
            min_distance=0.01,
            search_window=5,
            precision_limit=10 ** (-10),
            num_iter_max=10,
        )
        print(x_pos, y_pos, "test")
        x_pos = x_pos[:2]
        y_pos = y_pos[:2]
        kwargs_init = [
            {
                "theta_E": 1,
                "gamma": gamma,
                "e1": e1,
                "e2": e2,
                "center_x": 0,
                "center_y": 0,
            }
        ]
        kwargs_out_center, precision = solver_spep_center.constraint_lensmodel(
            x_pos, y_pos, kwargs_init
        )

        kwargs_init = [
            {
                "theta_E": 1,
                "gamma": gamma,
                "e1": 0,
                "e2": 0,
                "center_x": 0.1,
                "center_y": -0.1,
            }
        ]
        kwargs_out_ellipse, precision = solver_spep_ellipse.constraint_lensmodel(
            x_pos, y_pos, kwargs_init
        )

        npt.assert_almost_equal(
            kwargs_out_center[0]["center_x"], kwargs_lens[0]["center_x"], decimal=3
        )
        npt.assert_almost_equal(
            kwargs_out_center[0]["center_y"], kwargs_lens[0]["center_y"], decimal=3
        )
        npt.assert_almost_equal(kwargs_out_center[0]["center_y"], -0.1, decimal=3)

        npt.assert_almost_equal(
            kwargs_out_ellipse[0]["e1"], kwargs_lens[0]["e1"], decimal=3
        )
        npt.assert_almost_equal(
            kwargs_out_ellipse[0]["e2"], kwargs_lens[0]["e2"], decimal=3
        )
        npt.assert_almost_equal(kwargs_out_ellipse[0]["e1"], e1, decimal=3)

    def test_all_nfw(self):
        lensModel = LensModel(["SPEP"])
        solver_nfw_ellipse = Solver2Point(lensModel, solver_type="ELLIPSE")
        solver_nfw_center = Solver2Point(lensModel, solver_type="CENTER")
        spep = LensModel(["SPEP"])

        image_position_nfw = LensEquationSolver(LensModel(["SPEP", "NFW"]))
        sourcePos_x = 0.1
        sourcePos_y = 0.03
        deltapix = 0.05
        numPix = 100
        gamma = 1.9
        Rs = 0.1
        nfw = NFW()
        alpha_Rs = nfw.rho02alpha(1.0, Rs)
        phi_G, q = 0.5, 0.8
        e1, e2 = param_util.phi_q2_ellipticity(phi_G, q)
        kwargs_lens = [
            {
                "theta_E": 1.0,
                "gamma": gamma,
                "e1": e1,
                "e2": e2,
                "center_x": 0.1,
                "center_y": -0.1,
            },
            {"Rs": Rs, "alpha_Rs": alpha_Rs, "center_x": -0.5, "center_y": 0.5},
        ]
        x_pos, y_pos = image_position_nfw.findBrightImage(
            sourcePos_x,
            sourcePos_y,
            kwargs_lens,
            numImages=2,
            min_distance=deltapix,
            search_window=numPix * deltapix,
        )
        print(len(x_pos), "number of images")
        x_pos = x_pos[:2]
        y_pos = y_pos[:2]

        kwargs_init = [
            {
                "theta_E": 1,
                "gamma": gamma,
                "e1": e1,
                "e2": e2,
                "center_x": 0.0,
                "center_y": 0,
            },
            {"Rs": Rs, "alpha_Rs": alpha_Rs, "center_x": -0.5, "center_y": 0.5},
        ]
        kwargs_out_center, precision = solver_nfw_center.constraint_lensmodel(
            x_pos, y_pos, kwargs_init
        )
        source_x, source_y = spep.ray_shooting(x_pos[0], y_pos[0], kwargs_out_center)
        x_pos_new, y_pos_new = image_position_nfw.findBrightImage(
            source_x,
            source_y,
            kwargs_out_center,
            numImages=2,
            min_distance=deltapix,
            search_window=numPix * deltapix,
        )
        print(kwargs_out_center, "kwargs_out_center")
        npt.assert_almost_equal(x_pos_new[0], x_pos[0], decimal=2)
        npt.assert_almost_equal(y_pos_new[0], y_pos[0], decimal=2)

        npt.assert_almost_equal(
            kwargs_out_center[0]["center_x"], kwargs_lens[0]["center_x"], decimal=2
        )
        npt.assert_almost_equal(
            kwargs_out_center[0]["center_y"], kwargs_lens[0]["center_y"], decimal=2
        )
        npt.assert_almost_equal(kwargs_out_center[0]["center_y"], -0.1, decimal=2)

        kwargs_init = [
            {
                "theta_E": 1.0,
                "gamma": gamma,
                "e1": 0,
                "e2": 0,
                "center_x": 0.1,
                "center_y": -0.1,
            },
            {"Rs": Rs, "alpha_Rs": alpha_Rs, "center_x": -0.5, "center_y": 0.5},
        ]
        kwargs_out_ellipse, precision = solver_nfw_ellipse.constraint_lensmodel(
            x_pos, y_pos, kwargs_init
        )

        npt.assert_almost_equal(
            kwargs_out_ellipse[0]["e1"], kwargs_lens[0]["e1"], decimal=2
        )
        npt.assert_almost_equal(
            kwargs_out_ellipse[0]["e2"], kwargs_lens[0]["e2"], decimal=2
        )
        npt.assert_almost_equal(kwargs_out_ellipse[0]["e1"], e1, decimal=2)

    def test_all_spep_sis(self):
        lensModel = LensModel(["SPEP", "SIS"])
        solver_ellipse = Solver2Point(lensModel, solver_type="ELLIPSE")
        solver_center = Solver2Point(lensModel, solver_type="CENTER")
        spep = LensModel(["SPEP", "SIS"])
        image_position = LensEquationSolver(lensModel)
        sourcePos_x = 0.1
        sourcePos_y = 0.03
        deltapix = 0.05
        numPix = 100
        gamma = 1.9
        kwargs_lens = [
            {
                "theta_E": 1.0,
                "gamma": gamma,
                "e1": 0.2,
                "e2": -0.03,
                "center_x": 0.1,
                "center_y": -0.1,
            },
            {"theta_E": 0.6, "center_x": -0.5, "center_y": 0.5},
        ]
        x_pos, y_pos = image_position.findBrightImage(
            sourcePos_x,
            sourcePos_y,
            kwargs_lens,
            numImages=2,
            min_distance=deltapix,
            search_window=numPix * deltapix,
            precision_limit=10 ** (-10),
        )
        print(len(x_pos), "number of images")
        x_pos = x_pos[:2]
        y_pos = y_pos[:2]

        kwargs_init = [
            {
                "theta_E": 1,
                "gamma": gamma,
                "e1": 0.2,
                "e2": -0.03,
                "center_x": 0.0,
                "center_y": 0,
            },
            {"theta_E": 0.6, "center_x": -0.5, "center_y": 0.5},
        ]
        kwargs_out_center, precision = solver_center.constraint_lensmodel(
            x_pos, y_pos, kwargs_init
        )
        print(kwargs_out_center, "output")
        source_x, source_y = spep.ray_shooting(x_pos[0], y_pos[0], kwargs_out_center)
        x_pos_new, y_pos_new = image_position.findBrightImage(
            source_x,
            source_y,
            kwargs_out_center,
            numImages=2,
            min_distance=deltapix,
            search_window=numPix * deltapix,
        )
        npt.assert_almost_equal(x_pos_new[0], x_pos[0], decimal=3)
        npt.assert_almost_equal(y_pos_new[0], y_pos[0], decimal=3)

        npt.assert_almost_equal(
            kwargs_out_center[0]["center_x"], kwargs_lens[0]["center_x"], decimal=3
        )
        npt.assert_almost_equal(
            kwargs_out_center[0]["center_y"], kwargs_lens[0]["center_y"], decimal=3
        )
        npt.assert_almost_equal(kwargs_out_center[0]["center_y"], -0.1, decimal=3)

        kwargs_init = [
            {
                "theta_E": 1.0,
                "gamma": gamma,
                "e1": 0,
                "e2": 0,
                "center_x": 0.1,
                "center_y": -0.1,
            },
            {"theta_E": 0.6, "center_x": -0.5, "center_y": 0.5},
        ]
        kwargs_out_ellipse, precision = solver_ellipse.constraint_lensmodel(
            x_pos, y_pos, kwargs_init
        )

        npt.assert_almost_equal(
            kwargs_out_ellipse[0]["e1"], kwargs_lens[0]["e1"], decimal=3
        )
        npt.assert_almost_equal(
            kwargs_out_ellipse[0]["e2"], kwargs_lens[0]["e2"], decimal=3
        )
        npt.assert_almost_equal(kwargs_out_ellipse[0]["e1"], 0.2, decimal=3)

    def test_shapelet_cart(self):
        lens_model_list = ["SHAPELETS_CART", "SIS"]
        lens = LensModel(lens_model_list)
        solver = Solver2Point(lens, solver_type="SHAPELETS")
        image_position = LensEquationSolver(lens)
        sourcePos_x = 0.1
        sourcePos_y = 0.03
        deltapix = 0.05
        numPix = 100

        kwargs_lens = [
            {"coeffs": [1.0, 0.0, 0.1, 1.0], "beta": 1.0},
            {"theta_E": 1.0, "center_x": -0.1, "center_y": 0.1},
        ]
        x_pos, y_pos = image_position.findBrightImage(
            sourcePos_x,
            sourcePos_y,
            kwargs_lens,
            numImages=2,
            min_distance=deltapix,
            search_window=numPix * deltapix,
            precision_limit=10 ** (-10),
        )
        print(len(x_pos), "number of images")
        x_pos = x_pos[:2]
        y_pos = y_pos[:2]

        kwargs_init = [
            {"coeffs": [1.0, 0.0, 0.1, 1.0], "beta": 1.0},
            {"theta_E": 1.0, "center_x": -0.1, "center_y": 0.1},
        ]
        kwargs_out, precision = solver.constraint_lensmodel(x_pos, y_pos, kwargs_init)
        print(kwargs_out, "output")
        source_x, source_y = lens.ray_shooting(x_pos[0], y_pos[0], kwargs_out)
        x_pos_new, y_pos_new = image_position.findBrightImage(
            source_x,
            source_y,
            kwargs_out,
            numImages=2,
            min_distance=deltapix,
            search_window=numPix * deltapix,
        )
        npt.assert_almost_equal(x_pos_new[0], x_pos[0], decimal=3)
        npt.assert_almost_equal(y_pos_new[0], y_pos[0], decimal=3)

        npt.assert_almost_equal(
            kwargs_out[0]["coeffs"][1], kwargs_lens[0]["coeffs"][1], decimal=3
        )
        npt.assert_almost_equal(
            kwargs_out[0]["coeffs"][2], kwargs_lens[0]["coeffs"][2], decimal=3
        )

    def test_theta_E_phi(self):
        lensModel = LensModel(["SPEP", "SHEAR"])
        solver = Solver2Point(lensModel, solver_type="THETA_E_PHI")

        image_position = LensEquationSolver(lensModel)
        sourcePos_x = 0.1
        sourcePos_y = 0.03
        deltapix = 0.05
        numPix = 100
        gamma = 1.9
        kwargs_lens = [
            {
                "theta_E": 1.0,
                "gamma": gamma,
                "e1": 0.1,
                "e2": -0.03,
                "center_x": 0.1,
                "center_y": -0.1,
            },
            {"gamma1": 0.03, "gamma2": 0.0},
        ]
        x_pos, y_pos = image_position.findBrightImage(
            sourcePos_x,
            sourcePos_y,
            kwargs_lens,
            numImages=2,
            min_distance=deltapix,
            search_window=numPix * deltapix,
            precision_limit=10 ** (-15),
        )
        print(len(x_pos), "number of images")
        x_pos = x_pos[:2]
        y_pos = y_pos[:2]

        kwargs_init = [
            {
                "theta_E": 1.9,
                "gamma": gamma,
                "e1": 0.1,
                "e2": -0.03,
                "center_x": 0.1,
                "center_y": -0.1,
            },
            {"gamma1": 0.0, "gamma2": 0.03},
        ]
        kwargs_out, precision = solver.constraint_lensmodel(x_pos, y_pos, kwargs_init)
        print(kwargs_out, "output")
        source_x, source_y = lensModel.ray_shooting(x_pos[0], y_pos[0], kwargs_out)
        x_pos_new, y_pos_new = image_position.findBrightImage(
            source_x,
            source_y,
            kwargs_out,
            numImages=2,
            min_distance=deltapix,
            search_window=numPix * deltapix,
        )
        npt.assert_almost_equal(x_pos_new[0], x_pos[0], decimal=3)
        npt.assert_almost_equal(y_pos_new[0], y_pos[0], decimal=3)

        npt.assert_almost_equal(
            kwargs_out[0]["theta_E"], kwargs_lens[0]["theta_E"], decimal=3
        )
        npt.assert_almost_equal(
            kwargs_out[1]["gamma1"], kwargs_lens[1]["gamma1"], decimal=2
        )
        npt.assert_almost_equal(
            kwargs_out[1]["gamma2"], kwargs_lens[1]["gamma2"], decimal=2
        )

    def test_theta_E_ellipse(self):
        lensModel = LensModel(["SPEP", "SHEAR"])
        solver = Solver2Point(lensModel, solver_type="THETA_E_ELLIPSE")

        image_position = LensEquationSolver(lensModel)
        sourcePos_x = 0.1
        sourcePos_y = 0.03
        deltapix = 0.05
        numPix = 100
        gamma = 1.9
        kwargs_lens = [
            {
                "theta_E": 1.0,
                "gamma": gamma,
                "e1": 0.1,
                "e2": -0.03,
                "center_x": 0.1,
                "center_y": -0.1,
            },
            {"gamma1": 0.03, "gamma2": 0.0},
        ]
        x_pos, y_pos = image_position.findBrightImage(
            sourcePos_x,
            sourcePos_y,
            kwargs_lens,
            numImages=2,
            min_distance=deltapix,
            search_window=numPix * deltapix,
            precision_limit=10 ** (-15),
        )
        print(len(x_pos), "number of images")
        x_pos = x_pos[:2]
        y_pos = y_pos[:2]

        kwargs_init = [
            {
                "theta_E": 1.9,
                "gamma": gamma,
                "e1": -0.03,
                "e2": 0.1,
                "center_x": 0.1,
                "center_y": -0.1,
            },
            {"gamma1": 0.03, "gamma2": 0.0},
        ]
        kwargs_out, precision = solver.constraint_lensmodel(x_pos, y_pos, kwargs_init)
        print(kwargs_out, "output")
        source_x, source_y = lensModel.ray_shooting(x_pos[0], y_pos[0], kwargs_out)
        x_pos_new, y_pos_new = image_position.findBrightImage(
            source_x,
            source_y,
            kwargs_out,
            numImages=2,
            min_distance=deltapix,
            search_window=numPix * deltapix,
        )
        npt.assert_almost_equal(x_pos_new[0], x_pos[0], decimal=3)
        npt.assert_almost_equal(y_pos_new[0], y_pos[0], decimal=3)

        npt.assert_almost_equal(
            kwargs_out[0]["theta_E"], kwargs_lens[0]["theta_E"], decimal=3
        )
        npt.assert_almost_equal(kwargs_out[0]["e1"], kwargs_lens[0]["e1"], decimal=2)
        npt.assert_almost_equal(kwargs_out[0]["e2"], kwargs_lens[0]["e2"], decimal=2)

    def test_add_fixed_lens(self):
        lensModel = LensModel(lens_model_list=["SPEP", "SHEAR"])
        kwargs_lens_init = [{"theta_E": 1, "e2": 0}, {"gamma2": 0}]
        solver = Solver2Point(
            lensModel=lensModel, solver_type="THETA_E_PHI", decoupling=True
        )
        kwargs_fixed_added = solver.add_fixed_lens(
            kwargs_fixed_lens_list=[{}, {}], kwargs_lens_init=kwargs_lens_init
        )
        print(kwargs_fixed_added)
        assert kwargs_fixed_added[1]["gamma2"] == 0

        solver = Solver2Point(
            lensModel=lensModel, solver_type="THETA_E_ELLIPSE", decoupling=True
        )
        kwargs_fixed_added = solver.add_fixed_lens(
            kwargs_fixed_lens_list=[{}, {}], kwargs_lens_init=kwargs_lens_init
        )
        assert kwargs_fixed_added[0]["e2"] == 0

        lensModel = LensModel(lens_model_list=["SHAPELETS_CART"])
        solver = Solver2Point(
            lensModel=lensModel, solver_type="SHAPELETS", decoupling=True
        )
        kwargs_fixed_added = solver.add_fixed_lens(
            kwargs_fixed_lens_list=[{}, {}], kwargs_lens_init=kwargs_lens_init
        )
        assert len(kwargs_fixed_added) == len(kwargs_lens_init)


class TestRaise(unittest.TestCase):
    def test_raise(self):
        with self.assertRaises(ValueError):
            lensModel = LensModel(lens_model_list=["SPEP"])
            Solver2Point(lensModel=lensModel, solver_type="WRONG", decoupling=True)
        with self.assertRaises(ValueError):
            lensModel = LensModel(lens_model_list=["SPEP"])
            Solver2Point(lensModel=lensModel, solver_type="SHAPELETS", decoupling=True)
        with self.assertRaises(ValueError):
            lensModel = LensModel(lens_model_list=["SPEP", "SIS"])
            Solver2Point(
                lensModel=lensModel, solver_type="THETA_E_PHI", decoupling=True
            )
        with self.assertRaises(ValueError):
            lensModel = LensModel(lens_model_list=["SPEP", "SHEAR"])
            solver = Solver2Point(
                lensModel=lensModel, solver_type="THETA_E_PHI", decoupling=True
            )
            solver._solver_type = "WRONG"
            solver._update_kwargs(x=None, kwargs_list=None)
        with self.assertRaises(ValueError):
            lensModel = LensModel(lens_model_list=["SPEP", "SHEAR"])
            solver = Solver2Point(
                lensModel=lensModel, solver_type="THETA_E_PHI", decoupling=True
            )
            solver._solver_type = "WRONG"
            solver._extract_array(kwargs_list=None)
        with self.assertRaises(ValueError):
            lensModel = LensModel(lens_model_list=["SPEP", "SHEAR"])
            solver = Solver2Point(
                lensModel=lensModel, solver_type="THETA_E_PHI", decoupling=True
            )
            solver._solver_type = "WRONG"
            solver.add_fixed_lens(
                kwargs_fixed_lens_list=[None], kwargs_lens_init=[None]
            )


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