import numpy as np


def hydrogen_number_density_numpy(Delta_arr,zmean,cosmo_param,phy_const):
    omega_b = cosmo_param["omega_b"]
    h       = cosmo_param["hubble_param"]
    rho_c   = cosmo_param["rho_c_by_hsq"] * h**2.0 # cosmo_param[7] -->> rho_c_by_hsq
    Y       = cosmo_param["Y"]
    mu_arr  = MEAN_MOLECULAR_WEIGHT_NUMBA(Y)
    mu_H    = mu_arr[1]
    m_p     = phy_const["mass_proton_cgs"] #1.67e-24
    n_0     = rho_c * omega_b / (m_p * mu_H)
    n_H     = n_0 * (1 + zmean)**3.0 * Delta_arr
    return n_H

def MEAN_MOLECULAR_WEIGHT_NUMBA(Y,x_HII=1.0,x_HeII=0.0,x_HeIII=1.0):
    """
    mu_B         = mu_arr[0]
    mu_H         = mu_arr[1]
    mu_He        = mu_arr[2]
    mu_e         = mu_arr[3]
    nH_by_nB     = mu_arr[4]
    nHe_by_nB    = mu_arr[5]
    nHe_by_nH    = mu_arr[6]
    ne_by_nB     = mu_arr[7]
    ne_by_nH     = mu_arr[8]
    ne_by_nHe    = mu_arr[9]
    nH_by_ntotal = mu_arr[10]
    
    """
    # See the Notes_MMW.pdf for detailed description

    # Mean Molecular Weight of Baryons
    mu_B =  4.0 / (4.0 - 3.0*Y)

    # Mean Molecular Weight of Hydrogen
    mu_H =  1.0 / (1.0 - Y)

    # Mean Molecular Weight of Helium
    mu_He =  4.0 / Y

    # Assuming H and He are completely ionized for Default values
    X    = 1.0 - Y
    mu_e = 1.0 / (X * x_HII + (Y/4.0) * (x_HeII + 2.0*x_HeIII))


    # Ratio of number densities is inversaly proportional to mean 
    # molecular weight
    nH_by_nB     = mu_B / mu_H
    nHe_by_nB    = mu_B / mu_He
    nHe_by_nH    = mu_H / mu_He
    ne_by_nB     = mu_B / mu_e
    ne_by_nH     = mu_H / mu_e
    ne_by_nHe    = mu_He / mu_e
    nH_by_ntotal = 1.0 / (2.0 + 3.0 * nHe_by_nH)

    mu_arr = np.array([mu_B,mu_H,mu_He,mu_e,\
                       nH_by_nB,nHe_by_nB,\
                       nHe_by_nH,ne_by_nB,\
                       ne_by_nH,ne_by_nHe,\
                       nH_by_ntotal])
    return mu_arr
