__author__ = 'sibirrer'

from lenstronomy.LensModel.lens_model import LensModel
import lenstronomy.Util.param_util as param_util

import scipy.optimize
import numpy as np


class Solver2Point(object):
    """
    class to make the constraints for the solver
    """
    def __init__(self, lensModel, solver_type='CENTER', decoupling=True):
        self.lensModel = lensModel
        self._lens_mode_list = lensModel.lens_model_list
        self._solver_type = solver_type
        if lensModel.multi_plane or 'FOREGROUND_SHEAR' in self._lens_mode_list:
            self._decoupling = False
        else:
            self._decoupling = decoupling

    def constraint_lensmodel(self, x_pos, y_pos, kwargs_list):
        """

        :param x_pos: list of image positions (x-axis)
        :param y_pos: list of image position (y-axis)
        :param init: initial parameters
        :param kwargs_list: list of lens model kwargs
        :return: updated lens model that satisfies the lens equation for the point sources
        """
        init = self._extract_array(kwargs_list)
        if self._decoupling:
            alpha_0_x, alpha_0_y = self.lensModel.alpha(x_pos, y_pos, kwargs_list)
            alpha_1_x, alpha_1_y = self.lensModel.alpha(x_pos, y_pos, kwargs_list, k=0)
            x_sub = alpha_1_x - alpha_0_x
            y_sub = alpha_1_y - alpha_0_y
        else:
            x_sub, y_sub = np.zeros(2), np.zeros(2)
        a = self._subtract_constraint(x_sub, y_sub)
        x = self.solve(x_pos, y_pos, init, kwargs_list, a)
        kwargs_list = self._update_kwargs(x, kwargs_list)
        return kwargs_list

    def solve(self, x_pos, y_pos, init, kwargs_list, a):
        x = scipy.optimize.fsolve(self._F, init, args=(x_pos, y_pos, kwargs_list, a), xtol=1.49012e-10)#, factor=0.1)
        return x

    def _F(self, x, x_pos, y_pos, kwargs_list, a=0):
        kwargs_list = self._update_kwargs(x, kwargs_list)
        if self._decoupling:
            beta_x, beta_y = self.lensModel.ray_shooting(x_pos, y_pos, kwargs_list, k=0)
        else:
            beta_x, beta_y = self.lensModel.ray_shooting(x_pos, y_pos, kwargs_list)
        y = np.zeros(2)
        y[0] = beta_x[0] - beta_x[1]
        y[1] = beta_y[0] - beta_y[1]
        return y - a

    def _subtract_constraint(self, x_sub, y_sub):
        """

        :param x_pos:
        :param y_pos:
        :param x_sub:
        :param y_sub:
        :return:
        """
        a = np.zeros(2)
        a[0] = - x_sub[0] + x_sub[1]
        a[1] = - y_sub[0] + y_sub[1]
        return a

    def _update_kwargs(self, x, kwargs_list):
        """

        :param x: list of parameters corresponding to the free parameter of the first lens model in the list
        :param kwargs_list: list of lens model kwargs
        :return: updated kwargs_list
        """
        lens_model = self._lens_mode_list[0]
        if lens_model in ['SPEP', 'SPEMD', 'SIE', 'COMPOSITE', 'NFW_ELLIPSE']:
            if self._solver_type == 'CENTER':
                [center_x, center_y] = x
                kwargs_list[0]['center_x'] = center_x
                kwargs_list[0]['center_y'] = center_y
            elif self._solver_type == 'ELLIPSE':
                [e1, e2] = x
                phi_G, q = param_util.elliptisity2phi_q(e1, e2)
                kwargs_list[0]['q'] = q
                kwargs_list[0]['phi_G'] = phi_G

        elif lens_model in ['SHAPELETS_CART']:
            [c10, c01] = x
            coeffs = list(kwargs_list[0]['coeffs'])
            coeffs[1: 3] = [c10, c01]
            kwargs_list[0]['coeffs'] = coeffs
        else:
            raise ValueError("Lens model %s not supported for 2-point solver!" % lens_model)
        return kwargs_list

    def _extract_array(self, kwargs_list):
        """
        inverse of _update_kwargs
        :param kwargs_list:
        :return:
        """
        lens_model = self._lens_mode_list[0]
        if lens_model in ['SPEP', 'SPEMD', 'SIE', 'COMPOSITE', 'NFW_ELLIPSE']:
            if self._solver_type == 'CENTER':
                center_x = kwargs_list[0]['center_x']
                center_y = kwargs_list[0]['center_y']
                x = [center_x, center_y]
            elif self._solver_type == 'ELLIPSE':
                q = kwargs_list[0]['q']
                phi_G = kwargs_list[0]['phi_G']
                e1, e2 = param_util.phi_q2_elliptisity(phi_G, q)
                x = [e1, e2]
            else:
                raise ValueError("Solver type %s not valid for lens model %s" % (self._solver_type, lens_model))
        elif lens_model in ['SHAPELETS_CART']:
            coeffs = list(kwargs_list[0]['coeffs'])
            [c10, c01] = coeffs[1: 3]
            x = [c10, c01]
        else:
            raise ValueError("Lens model %s not supported for 2-point solver!" % lens_model)
        return x



