__author__ = "sibirrer"

import pytest
import numpy.testing as npt
from lenstronomy.LensModel.lens_param import LensParam


class TestParam(object):
    def setup_method(self):
        self.lens_model_list = [
            "SPEP",
            "INTERPOL_SCALED",
            "SHAPELETS_CART",
            "MULTI_GAUSSIAN_KAPPA",
        ]
        self.kwargs = [
            {
                "theta_E": 1.0,
                "gamma": 2,
                "e1": 0,
                "e2": 0,
                "center_x": 0,
                "center_y": 0,
            },  # 'SPEP
            {
                "scale_factor": 1,
                "grid_interp_x": None,
                "grid_interp_y": None,
                "f_x": None,
                "f_y": None,
            },  # 'INTERPOL_SCALED'
            {
                "coeffs": [1, 1, 1, 1, 1, 1],
                "beta": 1.0,
                "center_x": 0,
                "center_y": 0,
            },  # 'SHAPELETS_CART'
            {
                "amp": [1, 2],
                "sigma": [0.5, 1],
                "center_x": 0,
                "center_y": 0,
                "scale_factor": 1,
            },  # 'MULTI_GAUSSIAN_KAPPA'
        ]
        self.kwargs_sigma = [
            {
                "theta_E_sigma": 1.0,
                "gamma_sigma": 2,
                "e1_sigma": 0.1,
                "e2_sigma": 0.1,
                "center_x_sigma": 0,
                "center_y_sigma": 0,
            },  # 'SPEP
            {"scale_factor_sigma": 1},  # 'INTERPOL_SCALED'
            {
                "coeffs_sigma": 0.1,
                "beta_sigma": 1.0,
                "center_x_sigma": 0,
                "center_y_sigma": 0,
            },  # 'SHAPELETS_CART'
            {
                "amp_sigma": [1, 1],
                "sigma_sigma": [1, 1],
                "center_x_sigma": 0,
                "center_y_sigma": 0,
                "scale_factor_sigma": 1,
            },
        ]
        self.kwargs_fixed = [
            {},
            {"grid_interp_x": None, "grid_interp_y": None, "f_x": None, "f_y": None},
            {},
            {"sigma": [1, 2]},
        ]
        self.kwargs_mean = []
        for i in range(len(self.lens_model_list)):
            kwargs_mean_k = self.kwargs[i].copy()
            kwargs_mean_k.update(self.kwargs_sigma[i])
            self.kwargs_mean.append(kwargs_mean_k)
        self.param = LensParam(
            lens_model_list=self.lens_model_list,
            kwargs_fixed=self.kwargs_fixed,
            num_images=2,
            solver_type="SHAPELETS",
            num_shapelet_lens=6,
        )
        self.param_fixed = LensParam(
            lens_model_list=self.lens_model_list,
            kwargs_fixed=self.kwargs,
            num_images=4,
            solver_type="NONE",
            num_shapelet_lens=6,
        )

    def test_get_setParams(self):
        print(self.kwargs, "kwargs")
        args = self.param.set_params(self.kwargs)
        print(args, "args")
        kwargs_new, _ = self.param.get_params(args, i=0)
        print(kwargs_new, "kwargs_new")
        args_new = self.param.set_params(kwargs_new)
        print(args_new, "args_new")
        for k in range(len(args)):
            npt.assert_almost_equal(args[k], args_new[k], decimal=8)

        args = self.param_fixed.set_params(self.kwargs)
        kwargs_new, _ = self.param_fixed.get_params(args, i=0)
        args_new = self.param_fixed.set_params(kwargs_new)
        for k in range(len(args)):
            npt.assert_almost_equal(args[k], args_new[k], decimal=8)

    def test_param_name_list(self):
        lens_model_list = [
            "SHIFT",
            "FLEXION",
            "SIS_TRUNCATED",
            "SERSIC",
            "SERSIC_ELLIPSE_POTENTIAL",
            "SERSIC_ELLIPSE_KAPPA",
            "PJAFFE",
            "PJAFFE_ELLIPSE",
            "HERNQUIST_ELLIPSE",
            "INTERPOL",
            "INTERPOL_SCALED",
            "SHAPELETS_POLAR",
            "DIPOLE",
            "GAUSSIAN_ELLIPSE_KAPPA",
            "SERSIC_ELLIPSE_KAPPA",
            "SERSIC_ELLIPSE_GAUSS_DEC",
            "NFW_ELLIPSE_GAUSS_DEC",
            "CTNFW_GAUSS_DEC",
            "GAUSSIAN_ELLIPSE_POTENTIAL",
            "MULTI_GAUSSIAN_KAPPA",
            "MULTI_GAUSSIAN_KAPPA_ELLIPSE",
        ]
        lensParam = LensParam(lens_model_list, kwargs_fixed=None)
        param_name_list = lensParam._param_name_list
        assert len(lens_model_list) == len(param_name_list)

    def test_num_params(self):
        num, list = self.param.num_param()
        assert num == 20

    def test_shapelet_solver(self):
        lens_model_list = ["SHAPELETS_CART"]
        lensParam = LensParam(
            lens_model_list,
            kwargs_fixed=[{}],
            num_images=2,
            solver_type="SHAPELETS",
            num_shapelet_lens=8,
        )
        kwargs_lens = [
            {
                "beta": 1,
                "coeffs": [0, 1, 2, 3, 4, 5, 5, 7],
                "center_x": 0,
                "center_y": 0,
            }
        ]
        args = lensParam.set_params(kwargs_lens)
        kwargs_out, i = lensParam.get_params(args, i=0)
        assert kwargs_out[0]["coeffs"][1] == 0
        assert kwargs_out[0]["beta"] == kwargs_lens[0]["beta"]
        num, param_list = lensParam.num_param()
        assert num == 8

        lensParam = LensParam(
            lens_model_list,
            kwargs_fixed=[{}],
            num_images=4,
            solver_type="SHAPELETS",
            num_shapelet_lens=8,
        )
        kwargs_lens = [
            {
                "beta": 1,
                "coeffs": [0, 1, 2, 3, 4, 5, 5, 7],
                "center_x": 0,
                "center_y": 0,
            }
        ]
        args = lensParam.set_params(kwargs_lens)
        kwargs_out, i = lensParam.get_params(args, i=0)
        assert kwargs_out[0]["coeffs"][5] == 0
        assert kwargs_out[0]["beta"] == kwargs_lens[0]["beta"]
        num, param_list = lensParam.num_param()
        assert num == 5


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