from lenstronomy.LensModel.Solver.solver2point import Solver2Point
from lenstronomy.LensModel.Solver.solver4point import Solver4Point
import numpy as np

__all__ = ['Solver']


class Solver(object):
    """
    joint solve class to manage with type of solver to be executed and checks whether the requirements are fulfilled.

    """

    def __init__(self, solver_type, lensModel, num_images):
        """

        :param solver_type: string, option for specific solver type
         see detailed instruction of the Solver4Point and Solver2Point classes
        :param lensModel: instance of a LensModel() class
        :param num_images: int, number of images to be solved for
        """
        self._num_images = num_images
        self._lensModel = lensModel
        if self._num_images == 4:
            self._solver = Solver4Point(lensModel, solver_type=solver_type)
        elif self. _num_images == 2:
            self._solver = Solver2Point(lensModel, solver_type=solver_type)
        else:
            raise ValueError("%s number of images is not valid. Use 2 or 4!" % self._num_images)

    def constraint_lensmodel(self, x_pos, y_pos, kwargs_list, xtol=1.49012e-12):
        """

        :param x_pos: x-position constraints on images
        :param y_pos: y-position constraints on images
        :param kwargs_list: lens model keyword argument list
        :param xtol: tolerance level of solution when to stop the non-linear solver
        :return: updated lens model that satisfies the lens equation for the point sources, accuracy
        """
        return self._solver.constraint_lensmodel(x_pos, y_pos, kwargs_list, xtol=xtol)

    def update_solver(self, kwargs_lens, x_pos, y_pos):
        """


        :param kwargs_lens: lens model keyword argument list
        :param x_pos: x-position constraints on images
        :param y_pos: y-position constraints on images
        :return: updated lens model that satisfies the lens equation for the point sources
        """

        if not len(x_pos) == self._num_images:
            raise ValueError("Point source number %s must be as specified by the solver with number of images %s" %
                             (len(x_pos), self._num_images))
        kwargs_lens, precision = self.constraint_lensmodel(x_pos, y_pos, kwargs_lens)
        return kwargs_lens

    def check_solver(self, image_x, image_y, kwargs_lens):
        """
        returns the precision of the solver to match the image position

        :param kwargs_lens: full lens model (including solved parameters)
        :param image_x: point source in image
        :param image_y: point source in image
        :return: precision of Euclidean distances between the different rays arriving at the image positions
        """
        source_x, source_y = self._lensModel.ray_shooting(image_x, image_y, kwargs_lens)
        dist = np.sqrt((source_x - source_x[0]) ** 2 + (source_y - source_y[0]) ** 2)
        return dist

    def add_fixed_lens(self, kwargs_fixed_lens, kwargs_lens_init):
        """
        returns kwargs that are kept fixed during run, depending on options

        :param kwargs_fixed_lens: keyword argument list of fixed parameters (indicated by fitting argument of the user)
        :param kwargs_lens_init: Initial values of the full lens model keyword arguments
        :return: updated kwargs_fixed_lens, added fixed parameters being added (and replaced later on) by the
         non-linear solver.
        """
        kwargs_fixed_lens = self._solver.add_fixed_lens(kwargs_fixed_lens, kwargs_lens_init)
        return kwargs_fixed_lens
