# Created by Yen-Hsun Lin (Academia Sinica) in 04/2024.
# Copyright (c) 2024 Yen-Hsun Lin.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# any later version (see <http://www.gnu.org/licenses/>).
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.


__all__ = ['HaloSpike',
           'rhox',
           'M_sigma',
           'radiusInfluence',
           'radiusSchwarzschild',
           'nx',]


#---------- Import required utilities ----------#

from numpy import pi,broadcast_arrays,atleast_1d,zeros_like,nditer
from fractions import Fraction
from .constants import Constants,constant
from .params import params
from .sysmsg import FlagError


##########################################################################
#                                                                        #
#   General Classes and Functions for Numerics                           #
#                                                                        #
##########################################################################


"""
This script contains following classes and functions

Classes
------
1. HaloSpike

Functions
------
1. rhox
2. M_sigma
3. radiusInfluence
4. radiusSchwarzschild
5. nx
6. nxSpike

The docstrings should be sufficient for their self-explanations
"""


class HaloSpike(Constants):  
    """
    Class for constructing dark matter halo with spike due to supermassive black hole
    (SMBH) in the galactic center.

    /*-----------------------------------------------------------------------------*/

    
    ********************
    *                  *
    *   Class Inputs   *
    *                  *
    ********************
    
      mBH: SMBH mass, Msun
      tBH: SMBH age, years
    alpha: Slope of the spike profile


    ********************
    *                  *
    *    Attributes    *
    *                  *
    ********************

    When a HaloSpike instance is initialized, the following attributes will be assigned,

      mBH: SMBH mass, Msun
      tBH: SMBH age, years
    alpha: Slope of the spike profile
       rh: SMBH influence radius, kpc
      Rsp: Spike radius, kpc
    
    where the first four are the inputs with the last one auto-generated when initializing
    the instance. 
    
    >>> nx = HaloSpike(mBH=1e7,tBH=1e10,alpha='3/2')   # initializing HaloSpike instance
    
    Basic properties of the instance can be viewd by

    >>> nx
           SMBH mass: 1.000e+07 M_sun
         Spike slope: 3/2
       Initial slope: 1.000e+00
        Spike radius: 9.504e-10 kpc
    Influence radius: 3.412e-03 kpc

    In principle, alpha will be displayed in fractional form instead of str type as the input.
    Suppose one inputed alpha = '3/2' during the initialization, for such instance we can find
    that
    
    >>> nx.alpha                                               # get the attribute alpha
    Fraction(3,2)

    It can be converted into float number by

    >>> float(nx.alpha)
    1.5

    One can also reset alpha by

    >>> nx.alpha = '7/3'   # reset alpha
    >>> nx.alpha
    Fraction(7,3)
    
    The SMBH influence radius rh, kpc, is auto-generated by the stellar dispersion relation
    via M-sigma relation given the input SMBH mass. It can be viewed by

    >>> nx.rh   # kpc
    0.003411783398329804

    If you do not satisfy the generated rh, you can change it to your preferred value
     
    >>> nx.rh = 2.54e-3
    >>> nx.rh
    2.54e-03
    
    But this new value may not agree with the astrophysical consensus and you better have good
    reason to do so! However, if you want to restore the default rh generated by the stellar
    dispersion relation, you can simply do

    >>> nx.rh = None     # restore the default rh
    >>> nx.rh            # the default value
    0.003411783398329804
    

    **********************
    *                    *
    *      Callable      *
    *                    *
    **********************

    After initializing the HaloSpike instance, it is callable like normal function with the
    following required inputs. The output is the DM number density, #/cm^3, at r

        r: Distance to GC, kpc
       mx: DM mass, MeV
     sigv: DM annihilation cross section in the unit of cm^3/s
           None means no annihilation
     rhos: Charasteristic density, MeV/cm^3
       rs: Characteristic radius, kpc
        n: Slope of the DM profile outside R_spike

    The parameters (r,rhos,rs,n) are the same in the function snorer.rhox and they are the
    general parameterization for NFW profile. See our paper for more explanation.

    Now we can get the DM number density by

    >>> nx(r=1e-7,mx=1e-3,sigv=None,rhos=184,rs=24,n=2)
    7185012798674088.0

    See Phys. Rev. Lett. 83, 1719 (1999), arXiv:astro-ph/9906391 for the theoretical foundation
        JCAP 06, 004 (2023), arXiv:2301.08756 for the parameterizations deployed here
    """
    def __init__(self,mBH,tBH,alpha):
        self.mBH = mBH
        self.tBH = tBH
        self.alpha = alpha

        # Evaluated quantities
        self._rh = radiusInfluence(mBH)
        self._Rsp = radiusSchwarzschild(self.mBH)

    def __repr__(self):
        return '{:>18s}'.format('SMBH mass:') + ' {:>.3e} M_sun'.format(self.mBH) + '\n' +                 \
               '{:>18s}'.format('Spike slope:') + ' {}'.format(self._alpha) + '\n' + \
               '{:>18s}'.format('Spike radius:') + ' {:>.3e} kpc'.format(self.Rsp) + '\n' +              \
               '{:>18s}'.format('Influence radius:') + ' {:>.3e} kpc'.format(self.rh)
     
    @property
    def Rsp(self):
        return self._Rsp
    
    @property
    def rh(self):
        return self._rh
    
    @rh.setter
    def rh(self,new_rh):
        if new_rh is not None:
            self._rh = new_rh
        else:
            self._rh = radiusInfluence(self.mBH)

    @property
    def alpha(self):
        return self._alpha

    @alpha.setter
    def alpha(self,new_alpha):
        self._alpha = self._alpha_handling(new_alpha)

    # Convert input alpha from str type to fraction type
    def _alpha_handling(self,alpha):
        if alpha == '3/2':
            set_alpha = Fraction(3,2)
        elif alpha == '7/3':
            set_alpha = Fraction(7,3)
        else:
            raise FlagError('Spike index \'alpha\' is not recognized in the allowed list [\'3/2\', \'7/3\'].')
        return set_alpha  
    
    def _normN(self) -> float:
        """
        The normalization N
        
        Output
        ------
        normalization N
        """
        # ini parameters
        mBH = self.mBH
        rh = self.rh
        Rs = self.Rsp
        
        ri = 4*Rs
        alpha = float(self.alpha)
        
        def _fa(r,alpha):
            return (r**3/(3 - alpha) + 12*Rs*r**2/(alpha - 2) - 48*Rs**2*r/(alpha - 1) + 64*Rs**3/alpha)/r**alpha

        fh,fi = _fa(rh,alpha),_fa(ri,alpha)
        norm = mBH*self.Msun/4/pi/(fh - fi)
        return norm

    def _radiusSpike(self,rhos,rs) -> float:
        """
        Get the spike radius

        Input
        ------
        rhos: characteristic density, MeV/cm^3
        rs: characteristic radius, kpc

        Output
        ------
        R_spike: Spike radius, kpc
        """
        # ini parameter
        rh = self.rh
        alpha = float(self.alpha)
        N = self._normN()
        rhos = rhos*self.kpc2cm**3
        if self.alpha == Fraction(3,2):
            return (N/rhos/rs)**(3/4)*rh**(5/8)
        elif self.alpha == Fraction(7,3):
            return (N/rhos/rs)**alpha
    
    def _rhoPrime(self,r,rhos,rs) -> float:
        """
        Get rho'

        Input
        ------
        r: distance to GC, kpc
        rhos: characteristic density, MeV/cm^3
        rs: characteristic radius, kpc

        Output
        ------
        rhoPrime: MeV/cm^3
        """
        # ini parameters
        Rs = self.Rsp
        rh = self.rh
        alpha = float(self.alpha)
        
        ri = 4*Rs
        N = self._normN()
        Rsp = self._radiusSpike(rhos,rs)
        if self.alpha == Fraction(3,2):
            rhoN = N/rh**alpha
            rhoNp = rhoN*(rh/Rsp)**(7/3)
            if ri <= r < rh:
                rhoP = rhoN*(1 - ri/r)**3*(rh/r)**alpha
            else:
                rhoP = rhoNp*(Rsp/r)**(7/3)
            return rhoP/self.kpc2cm**3
        elif self.alpha == Fraction(7,3):
            rhoN = N/Rsp**(7/3)
            rhoP = rhoN*(1 - 4*Rs/r)**3*(Rsp/r)**alpha
            return rhoP/self.kpc2cm**3

    def _nxSpike(self,r,mx,sigv,rhos,rs,n) -> float:
        """
        DM number density with spike in the center

        Input
        ------
        r: distance to GC, kpc
        mx: DM mass, MeV
        sigv: DM annihilation cross section, in the unit of 1e-26 cm^3/s
            None indicates no annihilation
        rhos: The characteristic density, MeV/cm^3
        rs: The characteristic radius, kpc
        n: Slope of the DM profile

        Output
        ------
        number density: #/cm^3
        """
        # ini parameters
        Rs = self.Rsp
        tBH = self.tBH
        ri = 4*Rs
        
        if sigv is None:  # no DM annihilation
            Rsp = self._radiusSpike(rhos,rs)
            if r < ri:
                return 0
            elif ri <= r < Rsp:
                return self._rhoPrime(r,rhos,rs)/mx
            else:
                return rhox(r,rhos,rs,n)/mx
        else: # non-zero DM annihilation
            Rsp = self._radiusSpike(rhos,rs)
            sigv = sigv*1e-26
            rhoc = mx/sigv/tBH/self.year2Seconds
            if r < ri:
                return 0
            elif ri <= r < Rsp:
                rhoP = self._rhoPrime(r,rhos,rs)
                return rhoP*rhoc/(rhoP + rhoc)/mx
            else:
                rhoDM = rhox(r,rhos,rs,n)
                return rhoDM*rhoc/(rhoDM + rhoc)/mx
    
    def __call__(self,r,mx,sigv,rhos,rs,n):
        """
        Obtain the DM number density in the presence of spike
        
        Input
        ------
        r: distance to GC, kpc
        mx: DM mass, MeV
        sigv: DM annihilation cross section, in the unit of 1e-26 cm^3/s
            None indicates no annihilation
        rhos: The characteristic density, MeV/cm^3
        rs: The characteristic radius, kpc
        n: Slope of the DM profile

        Output
        ------
        number density: #/cm^3
        """
        return self._nxSpike(r,mx,sigv,rhos,rs,n)


def rhox(r,rhos,rs,n) -> float:
    """
    Dark matter density at r
    
    Parameters
    ----------
    r : array_like
        Distance from Earth to GC, kpc
    rhos : array_like
        Characteristic density, MeV/cm^3
    rs : array_like 
        Characteristic radius, kpc
    n : array_like
        Slope of the DM profile
    
    Returns
    -------
    out : scalar/ndarray
        DM density at r, MeV/cm^3. Out is scalar if all inputs
        are scalars

    /*---------------------------*/
    
    Additional Description
    ------

    This function evaluates the DM density profile

          rhox(r) = rhos/(r/rs)/(1+r/rs)^n

    where rhos and rs are characteristic density and radius respectively.
    wWhen (rhos,rs,n) = (184 MeV/cm^3, 24.42 kpc, 2), it is the famous
    Navarro–Frenk–White (NFW) profile. If divided by DM mass mx, it becomes
    the DM number density with unit #/cm^3

    See Phys. Rept. 405, 279 (2005), arXiv:hep-ph/0404175
    """
    rr = r/rs
    return rhos/(rr*(1 + rr)**n)


def M_sigma(mBH):
    """
    Stellar dispersion relation under the influence of supermassive black hole.
    Also known as M-sigma relation.

    Parameters
    ----------
    mBH : array_like
        SMBH mass, Msun

    Returns
    -------
    out : scalar/ndarray
        Stellar velocity dispersion, km/s. Out is scalar if the input is scalar
        too.

    /*---------------------------*/
    
    Additional Description
    ------

    This function evaluates the stellar dispersion sigma_s near SMBH

          log10(mBH/Msun) = 8.29 + 5.12*log10(sigma_s/200 km/s)

    See Nature 480, 215 (2011), arXiv:1112.1078
    """
    return 200*(mBH/1e8/1.9)**(1/5.1)


def radiusInfluence(mBH):
    """
    Influence radius of a supermassive black hole

    Parameters
    ----------
    mBH : array_like 
        Supermassive black hole mass, Msun
    
    Returns
    -------
    out : scalar/ndarray
        Influence radisu, kpc. Out is scalar if the input is scalar.

    /*---------------------------*/
    
    Additional Description
    ------

    This function evaluates the influence radius of a SMBH

          rh = G*mBH/sigma_s^2

    where sigma_s is the stellar dispersion near SMBH, also see
    the docstring in function M_sigma for extra information
    """
    sigma_s = M_sigma(mBH)
    rh = constant.G*mBH/sigma_s**2 # pc
    return rh*1e-3 # kpc


def radiusSchwarzschild(mBH) -> float:
    """
    Schawarzschild radius of a supermassive black hole

    Parameters
    ----------
    mBH : array_like
        Black hole mass, Msun

    Returns
    -------
    out : scalar/ndarray
        Schwarzschild radius, kpc. Out is scalar if the input is scalar.
    """
    mBH = mBH*constant.Msun_kg
    Rs = mBH*1.48e-25/constant.kpc2cm  # convert Rs into kpc
    return Rs


def _nx(r,mx,**kwargs):
    """
    Dark matter number density at distance r to the galactic center.
    Spike feature is not included.

    Parameters
    ----------
    r : array_like
        Distance to galactic center, kpc
    mx : array_like
        Dark matter mass, MeV
    **kwargs
        Keyword arguments for characteristic parameters
        of NFW profile.
    
    Returns
    -------
    out : scalar/ndarray
        Dark matter number density at r, 1/cm^3
    """
    rhos,rs,n = params.merge('halo',**kwargs).values()
    return rhox(r,rhos,rs,n)/mx


def _nxSpike(r,mx,**kwargs) -> float:
    """
    Dark matter number density at distance r to the galactic center.
    Spike feature is included.
    
    Parameters
    ----------
    r : array_like
        Distance to galactic center, kpc
    mx : array_like
        Dark matter mass, MeV
    **kwargs
        Keyword arguments for characteristic parameters of NFW profile and
        spike halo.
    
    Returns
    -------
    out : scalar/ndarray
        Dark matter number density at r with spike in the center, 1/cm^3
    """
    # --- Setup halo instance --- #
    # Get default configuration and update it with user inputs if exist
    rhos,rs,n,mBH,tBH,rh,alpha,sigv = params.merge('halo','spike',**kwargs).values()
    nx_spike = HaloSpike(mBH,tBH,alpha)
    # Use recorded rh instead of auto generated
    nx_spike.rh = rh

    # --- Prepare for vectorization --- #
    r = atleast_1d(r) # Let r be at least 1d array for easy manipulation
    R,MX = broadcast_arrays(r,mx)
    # Setup empty array to store nx values
    NX = zeros_like(R)
    # Use nditer to mimic vectorization and retrieve nx
    with nditer([R,MX,NX],op_flags=[['readonly'],['readonly'],['writeonly']]) as it:
        for r,m,nd in it:
            nd[...] = nx_spike(r,m,sigv,rhos,rs,n)
    # Return the result with dimension as the input
    return NX if NX.size > 1 else NX.item()

def nx(r,mx,is_spike=False,**kwargs):
    """
    Dark matter number density at distance r to the galactic center.
    
    Parameters
    ----------
    r : array_like
        Distance to galactic center, kpc
    mx : array_like
        Dark matter mass, MeV
    is_spike : bool
        Is spike feature included? Default is False.
    **kwargs
        Keyword arguments for characteristic parameters of NFW profile and
        spike halo. If 'is_spike = False', the parameters for configuring
        spiky halo will be deactivated. Default values assume Milky Way.
    
    Returns
    -------
    out : scalar/ndarray
        Dark matter number density at r with spike in the center, 1/cm^3
    """
    if is_spike is True:
        return _nxSpike(r,mx,**kwargs)
    elif is_spike is False:
        return _nx(r,mx,**kwargs)
    else:
        raise FlagError('\'is_spike\' must be a bool.')