__author__ = "sibirrer"

import lenstronomy.Util.param_util as param_util
from lenstronomy.LensModel.Profiles.base_profile import LensProfileBase
from lenstronomy.LensModel.Profiles.convergence import Convergence
import numpy as np

__all__ = ["Shear", "ShearGammaPsi", "ShearReduced"]


class Shear(LensProfileBase):
    """Class for external shear gamma1, gamma2 expression."""

    param_names = ["gamma1", "gamma2", "ra_0", "dec_0"]
    lower_limit_default = {"gamma1": -0.5, "gamma2": -0.5, "ra_0": -100, "dec_0": -100}
    upper_limit_default = {"gamma1": 0.5, "gamma2": 0.5, "ra_0": 100, "dec_0": 100}

    def function(self, x, y, gamma1, gamma2, ra_0=0, dec_0=0):
        """

        :param x: x-coordinate (angle)
        :param y: y0-coordinate (angle)
        :param gamma1: shear component
        :param gamma2: shear component
        :param ra_0: x/ra position where shear deflection is 0
        :param dec_0: y/dec position where shear deflection is 0
        :return: lensing potential
        """
        x_ = x - ra_0
        y_ = y - dec_0
        f_ = 1 / 2.0 * (gamma1 * x_ * x_ + 2 * gamma2 * x_ * y_ - gamma1 * y_ * y_)
        return f_

    def derivatives(self, x, y, gamma1, gamma2, ra_0=0, dec_0=0):
        """

        :param x: x-coordinate (angle)
        :param y: y0-coordinate (angle)
        :param gamma1: shear component
        :param gamma2: shear component
        :param ra_0: x/ra position where shear deflection is 0
        :param dec_0: y/dec position where shear deflection is 0
        :return: deflection angles
        """
        x_ = x - ra_0
        y_ = y - dec_0
        f_x = gamma1 * x_ + gamma2 * y_
        f_y = +gamma2 * x_ - gamma1 * y_
        return f_x, f_y

    def hessian(self, x, y, gamma1, gamma2, ra_0=0, dec_0=0):
        """

        :param x: x-coordinate (angle)
        :param y: y0-coordinate (angle)
        :param gamma1: shear component
        :param gamma2: shear component
        :param ra_0: x/ra position where shear deflection is 0
        :param dec_0: y/dec position where shear deflection is 0
        :return: f_xx, f_xy, f_yx, f_yy
        """
        gamma1 = gamma1
        gamma2 = gamma2
        kappa = 0
        f_xx = kappa + gamma1
        f_yy = kappa - gamma1
        f_xy = gamma2
        return f_xx, f_xy, f_xy, f_yy


class ShearGammaPsi(LensProfileBase):
    """
    class to model a shear field with shear strength and direction. The translation ot the cartesian shear distortions
    is as follow:

    .. math::
        \\gamma_1 = \\gamma_{ext} \\cos(2 \\phi_{ext})
        \\gamma_2 = \\gamma_{ext} \\sin(2 \\phi_{ext})

    """

    param_names = ["gamma_ext", "psi_ext", "ra_0", "dec_0"]
    lower_limit_default = {
        "gamma_ext": 0,
        "psi_ext": -np.pi,
        "ra_0": -100,
        "dec_0": -100,
    }
    upper_limit_default = {"gamma_ext": 1, "psi_ext": np.pi, "ra_0": 100, "dec_0": 100}

    def __init__(self):
        self._shear_e1e2 = Shear()
        super(ShearGammaPsi, self).__init__()

    @staticmethod
    def function(x, y, gamma_ext, psi_ext, ra_0=0, dec_0=0):
        """

        :param x: x-coordinate (angle)
        :param y: y0-coordinate (angle)
        :param gamma_ext: shear strength
        :param psi_ext: shear angle (radian)
        :param ra_0: x/ra position where shear deflection is 0
        :param dec_0: y/dec position where shear deflection is 0
        :return:
        """
        # change to polar coordinate
        r, phi = param_util.cart2polar(x - ra_0, y - dec_0)
        f_ = 1.0 / 2 * gamma_ext * r**2 * np.cos(2 * (phi - psi_ext))
        return f_

    def derivatives(self, x, y, gamma_ext, psi_ext, ra_0=0, dec_0=0):
        # rotation angle
        gamma1, gamma2 = param_util.shear_polar2cartesian(psi_ext, gamma_ext)
        return self._shear_e1e2.derivatives(x, y, gamma1, gamma2, ra_0, dec_0)

    def hessian(self, x, y, gamma_ext, psi_ext, ra_0=0, dec_0=0):
        gamma1, gamma2 = param_util.shear_polar2cartesian(psi_ext, gamma_ext)
        return self._shear_e1e2.hessian(x, y, gamma1, gamma2, ra_0, dec_0)


class ShearReduced(LensProfileBase):
    """Reduced shear distortions :math:`\\gamma' = \\gamma / (1-\\kappa)`. This
    distortion keeps the magnification as unity and, thus, does not change the size of
    apparent objects. To keep the magnification at unity, it requires.

    .. math::
        (1-\\kappa)^2) - \\gamma_1^2 - \\gamma_2^ = 1

    Thus, for given pair of reduced shear :math:`(\\gamma'_1, \\gamma'_2)`, an additional convergence term is calculated
    and added to the lensing distortions.
    """

    param_names = ["gamma1", "gamma2", "ra_0", "dec_0"]
    lower_limit_default = {"gamma1": -0.5, "gamma2": -0.5, "ra_0": -100, "dec_0": -100}
    upper_limit_default = {"gamma1": 0.5, "gamma2": 0.5, "ra_0": 100, "dec_0": 100}

    def __init__(self):
        self._shear = Shear()
        self._convergence = Convergence()
        super(ShearReduced, self).__init__()

    @staticmethod
    def _kappa_reduced(gamma1, gamma2):
        """Compute convergence such that magnification is unity.

        :param gamma1: reduced shear
        :param gamma2: reduced shear
        :return: kappa
        """
        kappa = 1 - 1.0 / np.sqrt(1 - gamma1**2 - gamma2**2)
        gamma1_ = (1 - kappa) * gamma1
        gamma2_ = (1 - kappa) * gamma2
        return kappa, gamma1_, gamma2_

    def function(self, x, y, gamma1, gamma2, ra_0=0, dec_0=0):
        """

        :param x: x-coordinate (angle)
        :param y: y0-coordinate (angle)
        :param gamma1: shear component
        :param gamma2: shear component
        :param ra_0: x/ra position where shear deflection is 0
        :param dec_0: y/dec position where shear deflection is 0
        :return: lensing potential
        """
        kappa, gamma1_, gamma2_ = self._kappa_reduced(gamma1, gamma2)
        f_shear = self._shear.function(x, y, gamma1_, gamma2_, ra_0, dec_0)
        f_kappa = self._convergence.function(x, y, kappa, ra_0, dec_0)
        return f_shear + f_kappa

    def derivatives(self, x, y, gamma1, gamma2, ra_0=0, dec_0=0):
        """

        :param x: x-coordinate (angle)
        :param y: y0-coordinate (angle)
        :param gamma1: shear component
        :param gamma2: shear component
        :param ra_0: x/ra position where shear deflection is 0
        :param dec_0: y/dec position where shear deflection is 0
        :return: deflection angles
        """
        kappa, gamma1_, gamma2_ = self._kappa_reduced(gamma1, gamma2)
        f_x_shear, f_y_shear = self._shear.derivatives(
            x, y, gamma1_, gamma2_, ra_0, dec_0
        )
        f_x_kappa, f_y_kappa = self._convergence.derivatives(x, y, kappa, ra_0, dec_0)
        return f_x_shear + f_x_kappa, f_y_shear + f_y_kappa

    def hessian(self, x, y, gamma1, gamma2, ra_0=0, dec_0=0):
        """

        :param x: x-coordinate (angle)
        :param y: y0-coordinate (angle)
        :param gamma1: shear component
        :param gamma2: shear component
        :param ra_0: x/ra position where shear deflection is 0
        :param dec_0: y/dec position where shear deflection is 0
        :return: f_xx, f_xy, f_yx, f_yy
        """
        kappa, gamma1_, gamma2_ = self._kappa_reduced(gamma1, gamma2)
        f_xx_g, f_xy_g, f_yx_g, f_yy_g = self._shear.hessian(
            x, y, gamma1_, gamma2_, ra_0, dec_0
        )
        f_xx_k, f_xy_k, f_yx_k, f_yy_k = self._convergence.hessian(
            x, y, kappa, ra_0, dec_0
        )
        f_xx = f_xx_g + f_xx_k
        f_yy = f_yy_g + f_yy_k
        f_xy = f_xy_g + f_xy_k
        return f_xx, f_xy, f_xy, f_yy
