# Class that implements the anisotropic spherical NFW DF with radially
# varying anisotropy of the Osipkov-Merritt type
import numpy

from ..potential import NFWPotential
from ..util import conversion
from .isotropicNFWdf import isotropicNFWdf
from .osipkovmerrittdf import _osipkovmerrittdf

_COEFFS = numpy.array(
    [
        -0.9958957901383353,
        4.2905266124525259,
        -7.6069046709185919,
        7.0313234865878806,
        -3.6920719890718035,
        0.8313023634615980,
        -0.2179687331774083,
        -0.0408426627412238,
        0.0802975743915827,
    ]
)


class osipkovmerrittNFWdf(_osipkovmerrittdf):
    """Class that implements the anisotropic spherical NFW DF with radially varying anisotropy of the Osipkov-Merritt type

    .. math::

        \\beta(r) = \\frac{1}{1+r_a^2/r^2}

    with :math:`r_a` the anistropy radius.

    """

    def __init__(self, pot=None, ra=1.4, rmax=1e4, ro=None, vo=None):
        """
        NAME:

            __init__

        PURPOSE:

            Initialize a NFW DF with Osipkov-Merritt anisotropy

        INPUT:

            pot - NFW potential which determines the DF

            ra - anisotropy radius (can be a Quantity)

           rmax= (1e4) maximum radius to consider (can be Quantity); set to numpy.inf to evaluate NFW w/o cut-off

           ro=, vo= galpy unit parameters

        OUTPUT:

            None

        HISTORY:

            2020-11-12 - Written - Bovy (UofT)
        """
        assert isinstance(pot, NFWPotential), "pot= must be potential.NFWPotential"
        _osipkovmerrittdf.__init__(self, pot=pot, ra=ra, rmax=rmax, ro=ro, vo=vo)
        self._Qtildemax = pot._amp / pot.a
        self._Qtildemin = -pot(self._rmax, 0, use_physical=False) / self._Qtildemax
        self._a2overra2 = self._pot.a**2.0 / self._ra2
        self._fQnorm = (
            self._a2overra2 / (4.0 * numpy.pi) / pot.a**1.5 / pot._amp**0.5
        )
        # Initialize isotropic version to use as part of fQ
        self._idf = isotropicNFWdf(pot=pot, rmax=rmax, ro=ro, vo=vo)

    def fQ(self, Q):
        """
        NAME:

            fQ

        PURPOSE

            Calculate the f(Q) portion of an Osipkov-Merritt NFW distribution function

        INPUT:

            Q - The Osipkov-Merritt 'energy' E-L^2/[2ra^2] (can be Quantity)

        OUTPUT:

            fQ - The value of the f(Q) portion of the DF

        HISTORY:

            2021-02-09 - Written - Bovy (UofT)

        """
        Qtilde = conversion.parse_energy(Q, vo=self._vo) / self._Qtildemax
        out = numpy.zeros_like(Qtilde)
        indx = (Qtilde > self._Qtildemin) * (Qtilde <= 1.0)
        # The 'ergodic' part
        out[indx] = self._idf.fE(-Q[indx])
        # The other part
        out[indx] += (
            self._fQnorm
            * numpy.polyval(_COEFFS, Qtilde[indx])
            / (
                Qtilde[indx] ** (2.0 / 3.0)
                * (numpy.log(Qtilde[indx]) / (1 - Qtilde[indx])) ** 2.0
            )
        )
        return out
