__author__ = "sibirrer"

import pytest
import unittest
import numpy as np
import numpy.testing as npt
from lenstronomy.PointSource.point_source_param import PointSourceParam


class TestParam(object):
    def setup_method(self):
        kwargs_fixed = [{}, {}, {}]
        num_point_sources_list = [4, 1, 1]
        fixed_magnification_list = [True, False, False]
        point_source_model_list = ["LENSED_POSITION", "SOURCE_POSITION", "UNLENSED"]
        self.param = PointSourceParam(
            model_list=point_source_model_list,
            kwargs_fixed=kwargs_fixed,
            num_point_source_list=num_point_sources_list,
            fixed_magnification_list=fixed_magnification_list,
        )
        self.kwargs = [
            {
                "ra_image": np.array([0, 0, 0, 0]),
                "dec_image": np.array([0, 0, 0, 0]),
                "source_amp": 1,
            },
            {"ra_source": 1, "dec_source": 1, "point_amp": np.array([1.0])},
            {"ra_image": [1], "dec_image": [1], "point_amp": np.array([1.0])},
        ]

        self.param_linear = PointSourceParam(
            model_list=point_source_model_list,
            kwargs_fixed=[{}, {}, {}],
            num_point_source_list=num_point_sources_list,
            linear_solver=False,
            fixed_magnification_list=fixed_magnification_list,
        )

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

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

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

        num, list = self.param_linear.num_param()
        assert num == 15

    def test_num_param_linear(self):
        num = self.param.num_param_linear()
        assert num == 3

        num = self.param_linear.num_param_linear()
        assert num == 0

    def test_init(self):
        ps_param = PointSourceParam(
            model_list=["UNLENSED"], kwargs_fixed=[{}], num_point_source_list=None
        )
        assert ps_param._num_point_sources_list[0] == 1


class TestRaise(unittest.TestCase):
    def test_raise(self):
        with self.assertRaises(ValueError):
            PointSourceParam(
                model_list=["BAD"],
                kwargs_fixed=[{}],
                kwargs_lower=None,
                kwargs_upper=[{"bla": 1}],
            )
        with self.assertRaises(ValueError):
            PointSourceParam(
                model_list=["BAD"],
                kwargs_fixed=[{}],
                kwargs_lower=[{"bla": 1}],
                kwargs_upper=None,
            )


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