__author__ = "sibirrer"

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

__all__ = ["SPEMD"]


class SPEMD(LensProfileBase):
    """Class for smooth power law ellipse mass density profile (SPEMD). This class
    effectively performs the FASTELL calculations by Renan Barkana. The parameters are
    changed and represent a spherically averaged Einstein radius an a logarithmic 3D
    mass profile slope.

    The SPEMD mass profile is defined as follow:

    .. math::
        \\kappa(x, y) = \\frac{3-\\gamma}{2} \\left(\\frac{\\theta_{E}}{\\sqrt{q x^2 + y^2/q + s^2}} \\right)^{\\gamma-1}

    with :math:`\\theta_{E}` is the (circularized) Einstein radius,
    :math:`\\gamma` is the negative power-law slope of the 3D mass distributions, :math:`q` is the minor/major axis ratio,
    and :math:`x` and :math:`y` are defined in a coordinate system aligned with the major and minor axis of the lens.

    the FASTELL definitions are as follows:

    The parameters are position :math:`(x1,x2)`, overall factor
    (:math:`b`), power (gam), axis ratio (arat) which is <=1, core radius
    squared (:math:`s2`), and the output potential (:math:`\\phi`).
    The projected mass density distribution, in units of the
    critical density, is

    .. math::
        \\kappa(x1,x2)=b_{fastell} \\left[u2+s2\\right]^{-gam},

    with :math:`u2=\\left[x1^2+x2^2/(arat^2)\\right]`.

    The conversion from lenstronomy definitions of this class to FASTELL are:

    .. math::
        q_{fastell} \\equiv q_{lenstronomy}

    .. math::
        gam \\equiv (\\gamma-1)/2

    .. math::
        b_{fastell} \\equiv (3-\\gamma)/2. * \\left(\\theta_{E}^2 / q\\right)^{gam}

    .. math::
        s2_{fastell} = s_{lenstronomy}^2 * q
    """

    param_names = ["theta_E", "gamma", "e1", "e2", "s_scale", "center_x", "center_y"]
    lower_limit_default = {
        "theta_E": 0,
        "gamma": 0,
        "e1": -0.5,
        "e2": -0.5,
        "s_scale": 0,
        "center_x": -100,
        "center_y": -100,
    }
    upper_limit_default = {
        "theta_E": 100,
        "gamma": 100,
        "e1": 0.5,
        "e2": 0.5,
        "s_scale": 100,
        "center_x": 100,
        "center_y": 100,
    }

    def __init__(self, suppress_fastell=False):
        """"""
        try:
            from fastell4py import fastell4py

            self._fastell4py_bool = True
            self.fastell4py = fastell4py
        except:
            self._fastell4py_bool = False
            if suppress_fastell:
                ImportWarning(
                    "module fastell4py not installed. You can get it from here: "
                    "https://github.com/sibirrer/fastell4py "
                    "Make sure you have a fortran compiler such that the installation works properly."
                )
                Warning(
                    "SPEMD model outputs are replaced by zeros as fastell4py package is not installed!"
                )
            else:
                raise ImportError(
                    "module fastell4py not installed. You can get it from here: "
                    "https://github.com/sibirrer/fastell4py "
                    "Make sure you have a fortran compiler such that the installation works properly."
                )
        super(SPEMD, self).__init__()

    def function(self, x, y, theta_E, gamma, e1, e2, s_scale, center_x=0, center_y=0):
        """

        :param x: x-coordinate (angle)
        :param y: y-coordinate (angle)
        :param theta_E: Einstein radius (angle), pay attention to specific definition!
        :param gamma: logarithmic slope of the power-law profile. gamma=2 corresponds to isothermal
        :param e1: eccentricity component
        :param e2: eccentricity component
        :param s_scale: smoothing scale in the center of the profile (angle)
        :param center_x: x-position of lens center
        :param center_y: y-position of lens center
        :return: lensing potential
        """
        x1, x2, q_fastell, gam, s2, q, phi_G = self.param_transform(
            x, y, theta_E, gamma, e1, e2, s_scale, center_x, center_y
        )
        compute_bool = self._parameter_constraints(q_fastell, gam, s2, q)
        if self._fastell4py_bool and self.is_not_empty(x1, x2) and compute_bool:
            potential = self.fastell4py.ellipphi(x1, x2, q_fastell, gam, arat=q, s2=s2)
            n = len(np.atleast_1d(x))
            if n <= 1:
                if np.shape(x) == ():
                    return np.array(potential[0])
        else:
            potential = np.zeros_like(x1)
        return potential

    def derivatives(
        self, x, y, theta_E, gamma, e1, e2, s_scale, center_x=0, center_y=0
    ):
        """

        :param x: x-coordinate (angle)
        :param y: y-coordinate (angle)
        :param theta_E: Einstein radius (angle), pay attention to specific definition!
        :param gamma: logarithmic slope of the power-law profile. gamma=2 corresponds to isothermal
        :param e1: eccentricity component
        :param e2: eccentricity component
        :param s_scale: smoothing scale in the center of the profile
        :param center_x: x-position of lens center
        :param center_y: y-position of lens center
        :return: deflection angles alpha_x, alpha_y
        """
        x1, x2, q_fastell, gam, s2, q, phi_G = self.param_transform(
            x, y, theta_E, gamma, e1, e2, s_scale, center_x, center_y
        )
        compute_bool = self._parameter_constraints(q_fastell, gam, s2, q)
        if self._fastell4py_bool and self.is_not_empty(x1, x2) and compute_bool:
            f_x_prim, f_y_prim = self.fastell4py.fastelldefl(
                x1, x2, q_fastell, gam, arat=q, s2=s2
            )
        else:
            f_x_prim, f_y_prim = np.zeros_like(x1), np.zeros_like(x1)
        cos_phi = np.cos(phi_G)
        sin_phi = np.sin(phi_G)

        f_x = cos_phi * f_x_prim - sin_phi * f_y_prim
        f_y = sin_phi * f_x_prim + cos_phi * f_y_prim
        return f_x, f_y

    def hessian(self, x, y, theta_E, gamma, e1, e2, s_scale, center_x=0, center_y=0):
        """

        :param x: x-coordinate (angle)
        :param y: y-coordinate (angle)
        :param theta_E: Einstein radius (angle), pay attention to specific definition!
        :param gamma: logarithmic slope of the power-law profile. gamma=2 corresponds to isothermal
        :param e1: eccentricity component
        :param e2: eccentricity component
        :param s_scale: smoothing scale in the center of the profile
        :param center_x: x-position of lens center
        :param center_y: y-position of lens center
        :return: Hessian components f_xx, f_xy, f_yx, f_yy
        """
        x1, x2, q_fastell, gam, s2, q, phi_G = self.param_transform(
            x, y, theta_E, gamma, e1, e2, s_scale, center_x, center_y
        )
        compute_bool = self._parameter_constraints(q_fastell, gam, s2, q)
        if self._fastell4py_bool and self.is_not_empty(x1, x2) and compute_bool:
            (
                f_x_prim,
                f_y_prim,
                f_xx_prim,
                f_yy_prim,
                f_xy_prim,
            ) = self.fastell4py.fastellmag(x1, x2, q_fastell, gam, arat=q, s2=s2)
            n = len(np.atleast_1d(x))
            if n <= 1:
                if np.shape(x) == ():
                    f_xx_prim, f_yy_prim, f_xy_prim = (
                        np.array(f_xx_prim[0]),
                        np.array(f_yy_prim[0]),
                        np.array(f_xy_prim[0]),
                    )
        else:
            f_xx_prim, f_yy_prim, f_xy_prim = (
                np.zeros_like(x1),
                np.zeros_like(x1),
                np.zeros_like(x1),
            )
        kappa = (f_xx_prim + f_yy_prim) / 2
        gamma1_value = (f_xx_prim - f_yy_prim) / 2
        gamma2_value = f_xy_prim

        gamma1 = np.cos(2 * phi_G) * gamma1_value - np.sin(2 * phi_G) * gamma2_value
        gamma2 = +np.sin(2 * phi_G) * gamma1_value + np.cos(2 * phi_G) * gamma2_value

        f_xx = kappa + gamma1
        f_yy = kappa - gamma1
        f_xy = gamma2
        return f_xx, f_xy, f_xy, f_yy

    def param_transform(
        self, x, y, theta_E, gamma, e1, e2, s_scale, center_x=0, center_y=0
    ):
        """Transforms parameters in the format of fastell4py.

        :param x: x-coordinate (angle)
        :param y: y-coordinate (angle)
        :param theta_E: Einstein radius (angle), pay attention to specific definition!
        :param gamma: logarithmic slope of the power-law profile. gamma=2 corresponds to
            isothermal
        :param e1: eccentricity component
        :param e2: eccentricity component
        :param s_scale: smoothing scale in the center of the profile
        :param center_x: x-position of lens center
        :param center_y: y-position of lens center
        :return: x-rotated, y-rotated, q_fastell, gam, s2, q, phi_G
        """
        phi_G, q = param_util.ellipticity2phi_q(e1, e2)
        x = np.array(x)
        y = np.array(y)
        x_shift = x - center_x
        y_shift = y - center_y
        q_fastell, gam, s2 = self.convert_params(theta_E, gamma, q, s_scale)
        cos_phi = np.cos(phi_G)
        sin_phi = np.sin(phi_G)

        x1 = cos_phi * x_shift + sin_phi * y_shift
        x2 = -sin_phi * x_shift + cos_phi * y_shift
        return x1, x2, q_fastell, gam, s2, q, phi_G

    @staticmethod
    def convert_params(theta_E, gamma, q, s_scale):
        """Converts parameter definitions into quantities used by the FASTELL fortran
        library.

        :param theta_E: Einstein radius
        :param gamma: 3D power-law slope of mass profile
        :param q: axis ratio minor/major
        :param s_scale: float, smoothing scale in the core
        :return: pre-factors to SPEMP profile for FASTELL
        """
        gam = (gamma - 1) / 2.0
        q_fastell = (3 - gamma) / 2.0 * (theta_E**2 / q) ** gam
        s2 = s_scale**2 * q
        return q_fastell, gam, s2

    @staticmethod
    def is_not_empty(x1, x2):
        """Check if float or not an empty array.

        :return: True if x1 and x2 are either floats/ints or an non-empty array, False
            if e.g. objects are []
        :rtype: bool
        """
        assert type(x1) == type(x2)

        if isinstance(x1, (list, tuple, np.ndarray)):
            if len(x1) != 0 and len(x2) != 0:
                return True
            else:
                return False
        else:
            return True

    @staticmethod
    def _parameter_constraints(q_fastell, gam, s2, q):
        """Sets bounds to parameters due to numerical stability.

        FASTELL has the following definitons:
        The parameters are position (x1,x2), overall factor
        (q), power (gam), axis ratio (arat) which is <=1, core radius
        squared (s2), and the output potential (phi).
        The projected mass density distribution, in units of the
        critical density, is kappa(x1,x2)=q [u2+s2]^(-gam), where
        u2=[x1^2+x2^2/(arat^2)].

        :param q_fastell: float, normalization of lens model, q_fastell = (3-gamma)/2. * (theta_E ** 2 / q) ** gam
        :param gam: float, slope parameter, gam = (gamma-1)/2.
        :param q: axis ratio
        :param s2: square of smoothing scale of the core
        :return: bool of whether or not to let the fastell provide to be evaluated or instead return zero(s)
        """
        if (
            q_fastell <= 0
            or q > 1
            or q < 0.01
            or gam > 0.999
            or gam < 0.001
            or not np.isfinite(q_fastell)
        ):
            return False
        return True
