import numpy as np
import mpmath
from scipy.integrate import quad
from astropy.cosmology import Planck18 as cosmo
import astropy.units as u
from astropy.units import cds
from astropy import constants as const


def constants():

    # units of this script are: 
    #       mass: solar masses
    #       time: giga years
    #   distance: mega parsecs
    
    Gu     = const.G.value          # in m^2 kg^-1 s^-2
    cu     = const.c.value          # in m /s 
    Msol_in_kg = const.M_sun.value  # in kg
    
    Gyr_in_s = 1.e9 * 1.*u.year.to(u.s)
    Mpc_in_m = 1.e6 * 1.*u.parsec.to(u.m)

    # Ho from Planck18, which we convert to units of 1/Gyrs
    H0_in_km_per_s_per_Mpc = cosmo.H(0).value
    H0_in_per_Gyr =  H0_in_km_per_s_per_Mpc * Gyr_in_s / (Mpc_in_m / 1000)

    # convert c into Mpc per Gyr
    c = cu * Gyr_in_s / Mpc_in_m
    
    # convert G to code units
    G = Gu * (Msol_in_kg*(Gyr_in_s**2.0)) / (Mpc_in_m**3.0)
    
    # one over 1 year in Gyrs^-1
    fgw = 1.0e9  # 

    # mass scaling in Msol
    Mscale = 1.0e7
    
    constants = (4.0*(G**(5.0/3.0))*(Mscale**(5.0/3.0))) \
                / (3.0*((np.pi)**(1.0/3.0))*(c**2.0)*(fgw**(4.0/3.0))*H0_in_per_Gyr)
	
    return (constants**(1.0/2.0))



def hmodel(theta,Mlow_choice,Mhigh_choice,zlow_choice,zhigh_choice):

    logno, beta, gamma, alpha, delta = theta

    no = 10.0**(logno)
	
    # Msol cancels from t = M/(10^7Msol)
    Mlow=float(Mlow_choice)  # 10.0**6.0 #Msol
    Mhigh=float(Mhigh_choice) #10.0**11.0 #Msol
    tlow=Mlow/(10.0**(7.0+delta))
    thigh=Mhigh/(10.0**(7.0+delta))

    Mincomplete_gamma_fns = ( mpmath.gammainc( ((5.0/3.0)-alpha), tlow )  - mpmath.gammainc( ((5.0/3.0)-alpha), thigh )  )
	

    # zintegral from 0 to 5
    OmegaM = cosmo.Om0
    OmegaL = cosmo.Ode0
    zlow=float(zlow_choice) #0.0 
    zhigh=float(zhigh_choice) #5.0


    zfn = lambda zi: (  ((1.0+zi)**(beta-(4.0/3.0)) * np.exp(-zi/gamma)) / ( (OmegaM*(1.0+zi)**3.0 + OmegaL)**(1.0/2.0) )  )
    Zint = quad(zfn,zlow,zhigh, epsabs=1.49e-12, epsrel=1.49e-12)

    h2 = no * ( 10.0**(delta*((5.0/3.0)-alpha)) ) * Mincomplete_gamma_fns * Zint[0]  * (1.0/np.log(10))


    h_no_constants = h2**(1.0/2.0)
    h = h_no_constants*constants()

    return h

