# 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',
           'dmNumberDensity',]


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

from numpy import pi
from fractions import Fraction
from .constants import Constants,constant
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. dmNumberDensity
6. dmNumberDensity_general

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


class HaloSpike(Constants):  
    """
    Class for constructing DM 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
    gamma: Slope of the initial 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
    
    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
        self.rh = radiusInfluence(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.Rs) + '\n' +              \
               '{:>18s}'.format('Influence radius:') + ' {:>.3e} kpc'.format(self.rh)
     
    @property
    def Rs(self):
        return radiusSchwarzschild(self.mBH)
    
    @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.Rs
        
        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.Rs
        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.Rs
        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:
    """
    DM density at given r
    
    Input
    ------
    r: distance to GC, kpc
    rhos: The characteristic density, MeV/cm^3
    rs: The characteristic radius, kpc
    n: Slope of the DM profile
    
    Output
    ------
    rhox: DM density at r, MeV/cm^3


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

    This function evaluates the DM density described by the profile

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

    where rhos and rs are characteristic density and radius respectively.
    when (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):
    """
    M-sigma relation for black holes

    Input
    -----
    mBH: SMBH mass, Msun

    Output
    -----
    sigma: stellar velocity dispersion, km/s


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

    This function evaluates the stellar dispersion sigma_s in SMBH adjacent

          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):
    """
    SMBH influence radius

    Input
    -----
    mBH: SMBH mass, Msun
    
    Output
    ------
    rh: influence radisu, kpc


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

    This function evaluates the influence radius of a SMBH

          rh = G*mBH/sigma_s^2

    where sigma_s is the stellar dispersion in SMBH adjacent, 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:
    """
    Calculating the Schawarzschild radius of a SMBH

    Input
    ------
    mBH: Black hole mass, Msun

    Output
    ------
    Rs: Schwarzschild radius, kpc
    """
    mBH = mBH*constant.Msun_kg
    Rs = mBH*1.48e-25/constant.kpc2cm  # convert Rs into kpc
    return Rs


def dmNumberDensity(r,mx,/,is_spike=True,sigv=None,tBH=1e10,profile='MW',alpha='3/2',**kwargs) -> float:
    """
    Obtain the DM number density at given r for MW or LMC
    
    In
    ------
    r: distance to GC, kpc
    mx: DM mass, MeV
    is_spike: Turn on/off spike feature, bool
    sigv: DM annihilation cross section, in the unit of 1e-26 cm^3/s
        None indicates no annihilation
    tBH: SMBH age, years
    profile: str, 'MW' or 'LMC'
    alpha: Slope of the spike, str type, '3/2' or '7/3'
    **kwargs: If you wish to have DM profile other than 'MW' or 'LMC',
        specify the desired rhos, rs, n, mBH or rh here. Those are not
        specified will be replaced by the values belong to the 'profile'
    
    Out
    ------
    number density: 1/cm^3

    See the docstrings in class haloSpike and function rhox for more detail
    """
    # setup profile parameters
    # they are stored in the dict with keys: {'rhos','rs','n','mBH',rh}
    if profile == 'MW':
        profile_params_dict = constant.MW_profile
    elif profile == 'LMC':
        profile_params_dict = constant.LMC_profile
    else:
        raise FlagError('Keyword argument \'profile\' must be either \'MW\' or \'LMC\'.')

    # if users have their own keyword arguments for the profile
    for key,value in kwargs.items():
        if profile_params_dict.get(key) is not None:  # check if the user-input exist or not
            profile_params_dict[key] = value
        else:
            raise TypeError(f'dmNumberDensity() got an unexpected keyword argument \'{key}\'')
    # assign values to variables
    rhos,rs,n,mBH,rh = profile_params_dict.values()

    if is_spike is True:
        nx = HaloSpike(mBH,tBH,alpha)
        if rh is None:   # use auto-generated rh
            return nx(r,mx,sigv,rhos,rs,n)
        else:
            nx.rh = rh   # use user-defined rh
            return nx(r,mx,sigv,rhos,rs,n)
    elif is_spike is False:       
        return rhox(r,rhos,rs,n)/mx
    else:
        raise FlagError('Keyword argument \'is_spike\' must be a boolean.')


# def dmNumberDensity(r,mx,is_spike=True,rh=None,sigv=None,tBH=1e10,profile='MW',alpha='3/2',gamma=1,**kwargs) -> float:
#     """
#     Obtain the DM number density at given r for MW or LMC
    
#     In
#     ------
#     r: distance to GC, kpc
#     mx: DM mass, MeV
#     is_spike: Turn on/off spike feature, bool
#     rh: SMBH influence radius, kpc
#         None indicates automatically calculated using the given mBH
#     sigv: DM annihilation cross section, in the unit of 1e-26 cm^3/s
#         None indicates no annihilation
#     tBH: SMBH age, years
#     profile: str, 'MW' or 'LMC'
#     alpha: Slope of the spike, str type, '3/2' or '7/3'
#     gamma: Slope of the initial profile
#     **kwargs: If you wish to have DM profile other than 'MW' or 'LMC',
#         specify the desired rhos, rs, n, mBH and rh here. Those not
#         specified will be replaced by the values belong to the 'profile'
    
#     Out
#     ------
#     number density: 1/cm^3

#     See the docstrings in class haloSpike and function rhox for more detail
#     """
#     if profile == 'MW':
#         rhos,rs,n,mBH,rh = 184,24.42,2,constant.M_SgrA,2e-3
#     elif profile == 'LMC':
#         rhos,rs,n,mBH,rh = 68,31.9,3,1e6,0.84e-3
#     else:
#         raise FlagError('Keyword argument \'profile\' must be either \'MW\' or \'LMC\'.')

#     if is_spike is True:
#         nx = HaloSpike(mBH,tBH,alpha,gamma)
#         if rh is None:
#             return nx(r,mx,sigv,rhos,rs,n)
#         else:
#             nx.rh = rh
#             return nx(r,mx,sigv,rhos,rs,n)
#     elif is_spike is False:       
#         return rhox(r,rhos,rs,n)/mx
#     else:
#         raise FlagError('Keyword argument \'is_spike\' must be a boolean.')


# def dmNumberDensity_general(r,mx,rhos,rs,n,mBH,is_spike=True,rh=None,sigv=None,tBH=1e10,alpha='3/2',gamma=1) -> float:
#     """
#     Obtain the DM number density at given r for user-defined DM halo profile
    
#     In
#     ------
#     r: distance to GC, kpc
#     mx: DM mass, MeV
#     rhos: The characteristic density, MeV/cm^3
#     rs: The characteristic radius, kpc
#     n: Slope of the DM profile
#     mBH: SMBH mass, Msun
#         if is_spike = False, mBH has no effect
#     is_spike: Turn on/off spike feature, bool
#     rh: SMBH influence radius, kpc
#         None indicates automatically calculated using the given mBH
#     sigv: DM annihilation cross section in the unit of 1e-26 cm^3/s, float
#         None indicates no annihilation
#     tBH: SMBH age, years
#     alpha: Slope of the spike, str type, '3/2' or '7/3'
#     gamma: Slope of the initial profile, float
    
#     Out
#     ------
#     number density: 1/cm^3

#     See the docstrings in class haloSpike and function rhox for more detail
#     """
#     if rh is None:
#         rh = radiusInfluence(mBH)

#     if is_spike is True:
#         nx = HaloSpike(mBH,tBH,alpha,gamma)
#         if rh is None:  # auto-calculated SMBH influence radius
#             return nx(r,mx,sigv,rhos,rs,n)
#         else:  # user-defined SMBH influence radius
#             nx.rh = rh
#             return nx(r,mx,sigv,rhos,rs,n)
#     elif is_spike is False:       
#         return rhox(r,rhos,rs,n)/mx
#     else:
#         raise FlagError('Keyword argument \'is_spike\' must be a boolean.')