"""Positive spherical density laws and numerically stable circular-speed kernels."""
import math
import numpy as np
from physical_units import GRAVITY_KPC_KMS2_PER_SOLAR_MASS as GRAVITY

def mass_kernel(kind,x):
    value=np.asarray(x,dtype=float)
    if np.any(~np.isfinite(value)) or np.any(value<0):raise ValueError('Invalid dimensionless radius')
    result=np.zeros_like(value);small=value<1e-2;large=~small
    if kind=='nfw':
        result[large]=np.log1p(value[large])-value[large]/(1+value[large])
        result[small]=sum((-1)**power*(power-1)/power*value[small]**power for power in range(2,14))
    elif kind=='core':
        result[large]=value[large]-np.arctan(value[large])
        result[small]=sum((-1)**power*value[small]**(2*power+3)/(2*power+3) for power in range(7))
    else:raise ValueError('Unknown halo kind')
    return result

def halo_speed_squared(radius,density,scale,kind):
    radius=np.asarray(radius,dtype=float)
    if not math.isfinite(density+scale) or density<=0 or scale<=0 or np.any(~np.isfinite(radius)) or np.any(radius<0):raise ValueError('Invalid halo parameter')
    x=radius/scale
    return 4*math.pi*GRAVITY*density*scale**2*np.divide(mass_kernel(kind,x),x,out=np.zeros_like(x),where=x>0)

def halo_density(radius,density,scale,kind):
    x=np.asarray(radius,dtype=float)/scale
    if kind=='nfw':return density/(x*(1+x)**2)
    if kind=='core':return density/(1+x*x)
    raise ValueError('Unknown halo kind')