__author__ = 'otto.hannuksela'


import numpy as np
from lenstronomy.LensModel.Profiles.base_profile import LensProfileBase

__all__ = ['PointMassShear']


class PointMassShear(LensProfileBase):
    """
    class to compute the physical deflection angle of a point mass with macro-model magnifications
    """
    # For simplicity, presume (X0, X1)=(0, 0) (centered at zero) and theta_E = 1.
    param_names = ['mu_0', 'mu_1']
    lower_limit_default = {'mu_0': -100, 'mu_1': -100}
    upper_limit_default = {'mu_0': 100,  'mu_1': 100 }

    def __init__(self):
        self.r_min = 10**(-25)
        super(PointMassShear, self).__init__()

    def function(self, x, y, mu_0, mu_1, theta_E=1., center_x=0, center_y=0):
        """

        :param x: x-coord (in angles)
        :param y: y-coord (in angles)
        :param theta_E: Einstein radius (in angles)
        :return: lensing potential
        """
        x_ = x - center_x
        y_ = y - center_y
        psi_macro_00_ = (mu_0-1.)/mu_0
        psi_macro_11_ = (mu_1-1.)/mu_1

        phi_macro     = 0.5*(psi_macro_00_*x**2 + psi_macro_11_*y**2)

        a = np.sqrt(x_**2 + y_**2)
        if isinstance(a, int) or isinstance(a, float):
            r = max(self.r_min, a)
        else:
            r = np.empty_like(a)
            r[a > self.r_min] = a[a > self.r_min]  #in the SIS regime
            r[a <= self.r_min] = self.r_min

        phi_point = theta_E**2*np.log(r)

        phi = phi_macro + phi_point # Add point + macro model together
        return phi

    def derivatives(self, x, y, mu_0, mu_1, theta_E=1., center_x=0, center_y=0):
        """

        :param x: x-coord (in angles)
        :param y: y-coord (in angles)
        :param theta_E: Einstein radius (in angles)
        :return: deflection angle (in angles)
        """
        x_ = x - center_x
        y_ = y - center_y
        psi_macro_00_ = (mu_0-1.)/mu_0
        psi_macro_11_ = (mu_1-1.)/mu_1

        a = np.sqrt(x_**2 + y_**2)
        if isinstance(a, int) or isinstance(a, float):
            r = max(self.r_min, a)
        else:
            r = np.empty_like(a)
            r[a > self.r_min] = a[a > self.r_min]  #in the SIS regime
            r[a <= self.r_min] = self.r_min
        alpha_macro_x = psi_macro_00_ * x_
        alpha_macro_y = psi_macro_00_ * y_
        alpha_point_x = theta_E**2 * x_ /r**2
        alpha_point_y = theta_E**2 * y_ /r**2
 
        alpha_x = alpha_macro_x + alpha_point_x
        alpha_y = alpha_macro_y + alpha_point_y

        return alpha_x, alpha_y

    def hessian(self, x, y, mu_0, mu_1, theta_E=1., center_x=0, center_y=0):
        """

        :param x: x-coord (in angles)
        :param y: y-coord (in angles)
        :param theta_E: Einstein radius (in angles)
        :return: hessian matrix (in angles)
        """
        x_ = x - center_x
        y_ = y - center_y

        psi_macro_00_ = (mu_0-1.)/mu_0
        psi_macro_11_ = (mu_1-1.)/mu_1

        C = theta_E**2
        a = x_**2 + y_**2
        if isinstance(a, int) or isinstance(a, float):
            r2 = max(self.r_min**2, a)
        else:
            r2 = np.empty_like(a)
            r2[a > self.r_min**2] = a[a > self.r_min**2]  #in the SIS regime
            r2[a <= self.r_min**2] = self.r_min**2
        f_point_xx = C * (y_**2-x_**2)/r2**2
        f_point_yy = C * (x_**2-y_**2)/r2**2
        f_point_xy = -C * 2*x_*y_/r2**2

        f_macro_xx = psi_macro_00_
        f_macro_yy = psi_macro_11_
        f_macro_xy = 0

        f_xx       = f_point_xx + f_macro_xx
        f_yy       = f_point_yy + f_macro_yy
        f_xy       = f_point_xy + f_macro_xy

        return f_xx, f_xy, f_xy, f_yy






__author__ = 'otto.hannuksela'


import numpy as np
from lenstronomy.LensModel.Profiles.base_profile import LensProfileBase

__all__ = ['PointMassShear']


class PointMassShear(LensProfileBase):
    """
    class to compute the physical deflection angle of a point mass with macro-model magnifications
    """
    # For simplicity, presume (X0, X1)=(0, 0) (centered at zero) and theta_E = 1.
    param_names = ['mu_0', 'mu_1']
    lower_limit_default = {'mu_0': -100, 'mu_1': -100}
    upper_limit_default = {'mu_0': 100,  'mu_1': 100 }

    def __init__(self):
        self.r_min = 10**(-25)
        super(PointMassShear, self).__init__()

    def function(self, x, y, mu_0, mu_1, theta_E=1., center_x=0, center_y=0):
        """

        :param x: x-coord (in angles)
        :param y: y-coord (in angles)
        :param theta_E: Einstein radius (in angles)
        :return: lensing potential
        """
        x_ = x - center_x
        y_ = y - center_y
        psi_macro_00_ = (mu_0-1.)/mu_0
        psi_macro_11_ = (mu_1-1.)/mu_1

        phi_macro     = 0.5*(psi_macro_00_*x**2 + psi_macro_11_*y**2)

        a = np.sqrt(x_**2 + y_**2)
        if isinstance(a, int) or isinstance(a, float):
            r = max(self.r_min, a)
        else:
            r = np.empty_like(a)
            r[a > self.r_min] = a[a > self.r_min]  #in the SIS regime
            r[a <= self.r_min] = self.r_min

        phi_point = theta_E**2*np.log(r)

        phi = phi_macro + phi_point # Add point + macro model together
        return phi

    def derivatives(self, x, y, mu_0, mu_1, theta_E=1., center_x=0, center_y=0):
        """

        :param x: x-coord (in angles)
        :param y: y-coord (in angles)
        :param theta_E: Einstein radius (in angles)
        :return: deflection angle (in angles)
        """
        x_ = x - center_x
        y_ = y - center_y
        psi_macro_00_ = (mu_0-1.)/mu_0
        psi_macro_11_ = (mu_1-1.)/mu_1

        a = np.sqrt(x_**2 + y_**2)
        if isinstance(a, int) or isinstance(a, float):
            r = max(self.r_min, a)
        else:
            r = np.empty_like(a)
            r[a > self.r_min] = a[a > self.r_min]  #in the SIS regime
            r[a <= self.r_min] = self.r_min
        alpha_macro_x = psi_macro_00_ * x_
        alpha_macro_y = psi_macro_00_ * y_
        alpha_point_x = theta_E**2 * x_ /r**2
        alpha_point_y = theta_E**2 * y_ /r**2
 
        alpha_x = alpha_macro_x + alpha_point_x
        alpha_y = alpha_macro_y + alpha_point_y

        return alpha_x, alpha_y

    def hessian(self, x, y, mu_0, mu_1, theta_E=1., center_x=0, center_y=0):
        """

        :param x: x-coord (in angles)
        :param y: y-coord (in angles)
        :param theta_E: Einstein radius (in angles)
        :return: hessian matrix (in angles)
        """
        x_ = x - center_x
        y_ = y - center_y

        psi_macro_00_ = (mu_0-1.)/mu_0
        psi_macro_11_ = (mu_1-1.)/mu_1

        C = theta_E**2
        a = x_**2 + y_**2
        if isinstance(a, int) or isinstance(a, float):
            r2 = max(self.r_min**2, a)
        else:
            r2 = np.empty_like(a)
            r2[a > self.r_min**2] = a[a > self.r_min**2]  #in the SIS regime
            r2[a <= self.r_min**2] = self.r_min**2
        f_point_xx = C * (y_**2-x_**2)/r2**2
        f_point_yy = C * (x_**2-y_**2)/r2**2
        f_point_xy = -C * 2*x_*y_/r2**2

        f_macro_xx = psi_macro_00_
        f_macro_yy = psi_macro_11_
        f_macro_xy = 0

        f_xx       = f_point_xx + f_macro_xx
        f_yy       = f_point_yy + f_macro_yy
        f_xy       = f_point_xy + f_macro_xy

        return f_xx, f_xy, f_xy, f_yy






class to compute the physical deflection angle of a point mass with macro-model magnifications # For simplicity, presume (X0, X1)=(0, 0) (centered at zero) and theta_E = 1. :param x: x-coord (in angles)         :param y: y-coord (in angles)         :param theta_E: Einstein radius (in angles)         :return: lensing potential #in the SIS regime # Add point + macro model together :param x: x-coord (in angles)         :param y: y-coord (in angles)         :param theta_E: Einstein radius (in angles)         :return: deflection angle (in angles) #in the SIS regime :param x: x-coord (in angles)         :param y: y-coord (in angles)         :param theta_E: Einstein radius (in angles)         :return: hessian matrix (in angles) #in the SIS regime