import numpy as np
from lenstronomy.LensModel.QuadOptimizer.multi_plane_fast import MultiplaneFast
from lenstronomy.LensModel.QuadOptimizer.param_manager import PowerLawFreeShear
from lenstronomy.LensModel.lens_model import LensModel
from lenstronomy.LensModel.Solver.lens_equation_solver import LensEquationSolver
import numpy.testing as npt
import pytest


class TestFastRayShooting(object):

    def setup(self):

        self.zlens, self.zsource = 0.5, 1.5
        epl_kwargs = {'theta_E': 0.8, 'center_x': 0.1, 'center_y': 0., 'e1': -0.2, 'e2': 0.1, 'gamma': 2.05}
        shear_kwargs = {'gamma1': 0.09, 'gamma2': -0.02}
        kwargs_macro = [epl_kwargs, shear_kwargs]

        self.x_image = np.array([0.65043538, -0.31109505, 0.78906059, -0.86222271])
        self.y_image = np.array([-0.89067493, 0.94851787, 0.52882605, -0.25403778])

        halo_list = ['SIS', 'SIS', 'SIS']
        halo_z = [self.zlens - 0.1, self.zlens, self.zlens + 0.4]
        halo_kwargs = [{'theta_E': 0.05, 'center_x': 0.3, 'center_y': -0.9},
                       {'theta_E': 0.01, 'center_x': 1.3, 'center_y': -0.5},
                       {'theta_E': 0.02, 'center_x': -0.4, 'center_y': -0.4}]

        self.kwargs_epl = kwargs_macro + halo_kwargs
        self.zlist_epl = [self.zlens, self.zlens] + halo_z
        self.lens_model_list_epl = ['EPL', 'SHEAR'] + halo_list

        self.lensModel = LensModel(self.lens_model_list_epl, self.zlens, self.zsource, self.zlist_epl,
                              multi_plane=True)

        self.param_class = PowerLawFreeShear(self.kwargs_epl)

    def test_rayshooting(self):

        solver = LensEquationSolver(self.lensModel)
        source_x, source_y = -0.05, -0.02
        x_image_true, y_image_true = solver.findBrightImage(source_x, source_y, self.kwargs_epl)

        fast_rayshooting = MultiplaneFast(x_image_true, y_image_true, self.zlens, self.zsource,
                                           self.lens_model_list_epl, self.zlist_epl,
                 astropy_instance=None, param_class=self.param_class, foreground_rays=None,
                 tol_source=1e-5, numerical_alpha_class=None)

        x_fore, y_fore, alpha_x_fore, alpha_y_fore = fast_rayshooting._ray_shooting_fast_foreground()
        xtrue, ytrue, alpha_xtrue, alpha_ytrue = self.lensModel.lens_model.\
            ray_shooting_partial(np.zeros_like(x_image_true), np.zeros_like(y_image_true), x_image_true, y_image_true, 0.,
                                 self.zlens, self.kwargs_epl)

        npt.assert_almost_equal(x_fore, xtrue)
        npt.assert_almost_equal(y_fore, ytrue)

        args_lens = self.param_class.kwargs_to_args(self.kwargs_epl)
        xfast, yfast = fast_rayshooting.ray_shooting_fast(args_lens)

        x, y = self.lensModel.ray_shooting(x_image_true, y_image_true, self.kwargs_epl)

        npt.assert_almost_equal(xfast, x)
        npt.assert_almost_equal(yfast, y)

        x_inner, y_inner = fast_rayshooting.lensModel.ray_shooting(x_image_true, y_image_true, self.kwargs_epl)
        npt.assert_almost_equal(x_inner, xfast)
        npt.assert_almost_equal(y_inner, yfast)

        foreground_rays = fast_rayshooting._foreground_rays

        fast_rayshooting_new = MultiplaneFast(x_image_true, y_image_true, self.zlens, self.zsource,
                                           self.lens_model_list_epl, self.zlist_epl,
                 astropy_instance=None, param_class=self.param_class, foreground_rays=foreground_rays,
                 tol_source=1e-5, numerical_alpha_class=None)

        xfast, yfast = fast_rayshooting_new.ray_shooting_fast(args_lens)

        npt.assert_almost_equal(xfast, x)
        npt.assert_almost_equal(yfast, y)

        chi_square_source = fast_rayshooting.source_plane_chi_square(args_lens)
        chi_square_total = fast_rayshooting.chi_square(args_lens)

        logL = fast_rayshooting.logL(args_lens)
        logL_true = -0.5 * chi_square_total
        npt.assert_almost_equal(logL, logL_true)
        npt.assert_almost_equal(chi_square_total, chi_square_source)

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