__author__ = 'dgilman'

# this file contains a class to compute the Navaro-Frenk-White profile
import numpy as np
from scipy.special import hyp2f1
from scipy.special import beta
from lenstronomy.LensModel.Profiles.base_profile import LensProfileBase

__all__ = ['GNFW']


class GNFW(LensProfileBase):
    """
    This class contains a double power law profile with flexible inner and outer logarithmic slopes g and n

    .. math::
        \\rho(r) = \\frac{\\rho_0}{r^{\\gamma}} \\frac{Rs^{n}}{\\left(r^2 + Rs^2 \\right)^{(n - \\gamma)/2}}

    For g = 1.0 and n=3, it is approximately the same as an NFW profile
    The original reference is [1]_.

    .. [1] Munoz, Kochanek and Keeton, (2001), astro-ph/0103009, doi:10.1086/322314

    TODO: implement the gravitational potential for this profile
    """
    profile_name = 'GNFW'
    param_names = ['Rs', 'alpha_Rs', 'center_x', 'center_y', 'gamma_inner', 'gamma_outer']
    lower_limit_default = {'Rs': 0, 'alpha_Rs': 0, 'center_x': -100, 'center_y': -100, 'gamma_inner': 0.1, 'gamma_outer': 1.0}
    upper_limit_default = {'Rs': 100, 'alpha_Rs': 10, 'center_x': 100, 'center_y': 100, 'gamma_inner': 2.9, 'gamma_outer': 10.0}

    def derivatives(self, x, y, Rs, alpha_Rs, gamma_inner, gamma_outer, center_x=0, center_y=0):
        """
        returns df/dx and df/dy of the function which are the deflection angles

        :param x: angular position (normally in units of arc seconds)
        :param y: angular position (normally in units of arc seconds)
        :param Rs: turn over point in the slope of the NFW profile in angular unit
        :param alpha_Rs: deflection (angular units) at projected Rs
        :param gamma_inner: logarithmic profile slope interior to Rs
        :param gamma_outer: logarithmic profile slope outside Rs
        :param center_x: center of halo (in angular units)
        :param center_y: center of halo (in angular units)
        :return: deflection angle in x, deflection angle in y
        """
        rho0_input = self.alpha2rho0(alpha_Rs, Rs, gamma_inner, gamma_outer)
        Rs = np.maximum(Rs, 0.00000001)
        x_ = x - center_x
        y_ = y - center_y
        R = np.sqrt(x_ ** 2 + y_ ** 2)
        f_x, f_y = self.nfwAlpha(R, Rs, rho0_input, gamma_inner, gamma_outer, x_, y_)
        return f_x, f_y

    def hessian(self, x, y, Rs, alpha_Rs, gamma_inner, gamma_outer, center_x=0, center_y=0):
        """

        :param x: angular position (normally in units of arc seconds)
        :param y: angular position (normally in units of arc seconds)
        :param Rs: turn over point in the slope of the NFW profile in angular unit
        :param alpha_Rs: deflection (angular units) at projected Rs
        :param gamma_inner: logarithmic profile slope interior to Rs
        :param gamma_outer: logarithmic profile slope outside Rs
        :param center_x: center of halo (in angular units)
        :param center_y: center of halo (in angular units)
        :return: Hessian matrix of function d^2f/dx^2, d^2/dxdy, d^2/dydx, d^f/dy^2
        """
        rho0_input = self.alpha2rho0(alpha_Rs, Rs, gamma_inner, gamma_outer)
        x_ = x - center_x
        y_ = y - center_y
        R = np.sqrt(x_ ** 2 + y_ ** 2)
        R = np.maximum(R, 0.00000001)
        kappa = self.density_2d(R, 0, Rs, rho0_input, gamma_inner, gamma_outer)
        gamma1, gamma2 = self.nfwGamma(R, Rs, rho0_input, gamma_inner, gamma_outer, x_, y_)
        f_xx = kappa + gamma1
        f_yy = kappa - gamma1
        f_xy = gamma2
        return f_xx, f_xy, f_xy, f_yy

    @staticmethod
    def density(R, Rs, rho0, gamma_inner, gamma_outer):
        """
        three dimensional NFW profile

        :param R: radius of interest
        :type Rs: scale radius
        :param rho0: central density normalization
        :param gamma_inner: logarithmic profile slope interior to Rs
        :param gamma_outer: logarithmic profile slope outside Rs
        :return: rho(R) density
        """
        x = R/Rs
        outer_slope = (gamma_outer-gamma_inner)/2
        return rho0 / (x**gamma_inner * (1 + x ** 2) ** outer_slope)

    def density_lens(self, r, Rs, alpha_Rs, gamma_inner, gamma_outer):
        """
        computes the density at 3d radius r given lens model parameterization.
        The integral in the LOS projection of this quantity results in the convergence quantity.

        :param r: 3d radios
        :param Rs: scale radius
        :param alpha_Rs: deflection at Rs
        :param gamma_inner: logarithmic profile slope interior to Rs
        :param gamma_outer: logarithmic profile slope outside Rs
        :return: density rho(r)
        """
        rho0 = self.alpha2rho0(alpha_Rs, Rs, gamma_inner, gamma_outer)
        return self.density(r, Rs, rho0, gamma_inner, gamma_outer)

    def density_2d(self, x, y, Rs, rho0, gamma_inner, gamma_outer, center_x=0, center_y=0):
        """
        projected two dimensional profile

        :param x: angular position (normally in units of arc seconds)
        :param y: angular position (normally in units of arc seconds)
        :param Rs: turn over point in the slope of the NFW profile in angular unit
        :param rho0: density normalization at Rs
        :param gamma_inner: logarithmic profile slope interior to Rs
        :param gamma_outer: logarithmic profile slope outside Rs
        :param center_x: profile center (same units as x)
        :param center_y: profile center (same units as x)
        :return: Epsilon(R) projected density at radius R
        """
        x_ = x - center_x
        y_ = y - center_y
        R = np.sqrt(x_ ** 2 + y_ ** 2)
        x = R / Rs
        Fx = self._f(x, gamma_inner, gamma_outer)
        return 2 * rho0 * Rs * Fx

    @staticmethod
    def mass_3d(r, Rs, rho0, gamma_inner, gamma_outer):
        """
        mass enclosed a 3d sphere or radius r

        :param r: 3d radius
        :param Rs: scale radius
        :param rho0: density normalization
        :param gamma_inner: logarithmic profile slope interior to Rs
        :param gamma_outer: logarithmic profile slope outside Rs
        :return: M(<r)
        """
        Rs = float(Rs)
        const = 4 * np.pi * r ** 3 * rho0 * (Rs/r) ** gamma_inner
        m_3d = const/(3-gamma_inner) * hyp2f1((3-gamma_inner)/2,
                               (gamma_outer-gamma_inner)/2,
                               (5-gamma_inner)/2, -(r/Rs)**2)
        return m_3d

    def mass_3d_lens(self, r, Rs, alpha_Rs, gamma_inner, gamma_outer):
        """
        mass enclosed a 3d sphere or radius r.
        This function takes as input the lensing parameterization.

        :param r: 3d radius
        :param Rs: scale radius
        :param alpha_Rs: deflection angle at Rs
        :param gamma_inner: logarithmic profile slope interior to Rs
        :param gamma_outer: logarithmic profile slope outside Rs
        :return: M(<r)
        """
        rho0 = self.alpha2rho0(alpha_Rs, Rs, gamma_inner, gamma_outer)
        m_3d = self.mass_3d(r, Rs, rho0, gamma_inner, gamma_outer)
        return m_3d

    def mass_2d(self, R, Rs, rho0, gamma_inner, gamma_outer):
        """
        mass enclosed a 2d cylinder or projected radius R

        :param R: 3d radius
        :param Rs: scale radius
        :param rho0: central density normalization
        :param gamma_inner: logarithmic profile slope interior to Rs
        :param gamma_outer: logarithmic profile slope outside Rs
        :return: mass in cylinder
        """
        R = np.maximum(R, 0.00000001)
        x = R / Rs
        gx = self._g(x, gamma_inner, gamma_outer)
        m_2d = 4 * rho0 * Rs * R ** 2 * gx / x ** 2 * np.pi
        return m_2d

    def nfwAlpha(self, R, Rs, rho0, gamma_inner, gamma_outer, ax_x, ax_y):
        """

        deflection angel of NFW profile (times Sigma_crit D_OL) along the projection to coordinate 'axis'

        :param R: 3d radius
        :param Rs: scale radius
        :param rho0: central density normalization
        :param gamma_inner: logarithmic profile slope interior to Rs
        :param gamma_outer: logarithmic profile slope outside Rs
        :param ax_x: x coordinate relative to center
        :param ax_y: y coordinate relative to center
        :return: Epsilon(R) projected density at radius R
        """
        R = np.maximum(R, 0.00000001)
        x = R / Rs
        gx = self._g(x, gamma_inner, gamma_outer)
        a = 4 * rho0 * Rs * R * gx / x ** 2 / R
        return a * ax_x, a * ax_y

    def nfwGamma(self, R, Rs, rho0, gamma_inner, gamma_outer, ax_x, ax_y):
        """

        shear gamma of NFW profile (times Sigma_crit) along the projection to coordinate 'axis'

        :param R: 3d radius
        :param Rs: scale radius
        :param rho0: central density normalization
        :param gamma_inner: logarithmic profile slope interior to Rs
        :param gamma_outer: logarithmic profile slope outside Rs
        :param ax_x: x coordinate relative to center
        :param ax_y: y coordinate relative to center
        :return: Epsilon(R) projected density at radius R
        """
        R = np.maximum(R, 0.00000001)
        x = R / Rs
        gx = self._g(x, gamma_inner, gamma_outer)
        Fx = self._f(x, gamma_inner, gamma_outer)
        a = 2 * rho0 * Rs * (2 * gx / x ** 2 - Fx)  # /x #2*rho0*Rs*(2*gx/x**2 - Fx)*axis/x
        return a * (ax_y ** 2 - ax_x ** 2) / R ** 2, -a * 2 * (ax_x * ax_y) / R ** 2

    @staticmethod
    def _f(X, g, n):
        """
        analytic solution of the projection integral

        :param X: R/Rs
        :type X: float >0
        :param g: logarithmic profile slope interior to Rs
        :param n: logarithmic profile slope exterior to Rs
        :return: solution to the projection integral
        """
        if n == 3:
            n = 3.001  # for numerical stability
        hyp2f1_term = hyp2f1((n-1)/2, g/2, n/2, 1/(1+X**2))
        beta_term = beta((n-1)/2, 0.5)
        return 0.5 * beta_term * hyp2f1_term * (1+X**2) ** ((1-n)/2)

    @staticmethod
    def _g(X, g, n):
        """
        analytic solution of integral for NFW profile to compute deflection angel and gamma

        :param X: R/Rs
        :type X: float >0
        :param g: logarithmic profile slope interior to Rs
        :param n: logarithmic profile slope exterior to Rs
        :return: solution of the integral over projected mass
        """
        if n == 3:
            n = 3.001  # for numerical stability
        xi = 1 + X**2
        hyp2f1_term = hyp2f1((n - 3) / 2, g / 2, n / 2, 1 / xi)
        beta_term_1 = beta((n - 3) / 2, (3-g)/2)
        beta_term_2 = beta((n-3)/2, 1.5)
        return 0.5 * (beta_term_1 - beta_term_2 * hyp2f1_term * xi ** ((3-n)/2))

    def alpha2rho0(self, alpha_Rs, Rs, gamma_inner, gamma_outer):

        """
        convert angle at Rs into rho0

        :param alpha_Rs: deflection angle at RS
        :param Rs: scale radius
        :param gamma_inner: logarithmic profile slope interior to Rs
        :param gamma_outer: logarithmic profile slope outside Rs
        :return: density normalization (characteristic density)
        """

        gx = self._g(1.0, gamma_inner, gamma_outer)
        rho0 = alpha_Rs / (4. * Rs ** 2 * gx / 1.0 ** 2)
        return rho0

    def rho02alpha(self, rho0, Rs, gamma_inner, gamma_outer):

        """
        convert rho0 to angle at Rs

        :param rho0: density normalization (characteristic density)
        :param Rs: scale radius
        :param gamma_inner: logarithmic profile slope interior to Rs
        :param gamma_outer: logarithmic profile slope outside Rs
        :return: deflection angle at RS
        """
        gx = self._g(1.0, gamma_inner, gamma_outer)
        alpha_Rs = rho0 * (4. * Rs ** 2 * gx / 1.0 ** 2)
        return alpha_Rs
