from argparse import ArgumentParser

def register_arguments_cosmo(parser:ArgumentParser):
    """
    Register the arguments for the cosmological parameters.
    """
    parser.add_argument("--h", type=float, default=0.6732, help="Hubble constant.")
    parser.add_argument("-Om","--Omega_m", type=float, default=0.302, help="Matter density parameter.")
    parser.add_argument("-Ob","--Omega_b", type=float, default=0.049, help="Baryon density parameter.")
    parser.add_argument("-Oq","--Omega_q", type=float, default=0.6842, help="Dark energy density parameter.")
    parser.add_argument("--n_s", type=float, default=0.968, help="Spectral index.")
    parser.add_argument("-s8","--sigma8", type=float, default=0.815, help="RMS amplitude of the density fluctuations.")
    parser.add_argument("--A_s", type=float, default=2.148752e-09, help="Amplitude of the primordial power spectrum.")
    parser.add_argument("--Tcmb", type=float, default=2.7255, help="CMB temperature.")
    parser.add_argument("--k_p", type=float, default=0.05, help="Pivot scale.")
    parser.add_argument("--N_ur", type=float, default=2.046, help="Effective number of ultra-relativistic species.")
    parser.add_argument("--m_nu1", type=float, default=0.06, help="Mass of the first neutrino species.")
    parser.add_argument("--m_nu2", type=float, default=0.0, help="Mass of the second neutrino species.")
    parser.add_argument("--m_nu3", type=float, default=0.0, help="Mass of the third neutrino species.")
    parser.add_argument("--w_0", type=float, default=-1.0, help="Dark energy equation of state parameter.")
    parser.add_argument("--w_a", type=float, default=0.0, help="Dark energy equation of state parameter.")
    parser.add_argument("--fnl", type=float, default=100.0, help="Local non-Gaussianity parameter.")
    parser.add_argument("--gnl", type=float, default=0.0, help="Equilateral non-Gaussianity parameter.")



def parse_arguments_cosmo(parsed_args):
    """
    Parse the arguments for the cosmological parameters.
    """
    cosmo_dict = dict(
            ## Common parameters
            h=parsed_args.h,
            Omega_m=parsed_args.Omega_m,
            Omega_b=parsed_args.Omega_b,
            Omega_q=parsed_args.Omega_q,
            Omega_k=0.0,
            n_s=parsed_args.n_s,
            sigma8=parsed_args.sigma8,
            w_0=parsed_args.w_0,
            w_a=parsed_args.w_a,
            A_s=parsed_args.A_s,

            ## monofonic ICs additional parameters
            Tcmb=parsed_args.Tcmb,
            k_p=parsed_args.k_p,
            N_ur=parsed_args.N_ur,
            m_nu1=parsed_args.m_nu1,
            m_nu2=parsed_args.m_nu2,
            m_nu3=parsed_args.m_nu3,
            fnl=parsed_args.fnl,
            gnl=parsed_args.gnl,

            ## sbmy ICs additional parameters
            Omega_r=0.0,
            k_max=10.0, 
            tau_reio=0.06,
            WhichSpectrum="EH", # EH or class
            w0_fld=parsed_args.w_0,
            wa_fld=parsed_args.w_a,
            )
    
    return cosmo_dict

def z2a(z):
    return 1.0/(1.0+z)


cosmo_defaults = {
        "h": 0.6732,
        "Omega_m": 0.302,
        "Omega_b": 0.049,
        "Omega_q": 0.6842,
        "Omega_k":0.0,
        "Omega_r": 0.0,
        "n_s": 0.968,
        "sigma8": 0.815,
        "A_s": 2.148752e-09,
        "Tcmb": 2.7255,
        "k_p": 0.05,
        "N_ur": 2.046,
        "m_nu1": 0.06,
        "m_nu2": 0.0,
        "m_nu3": 0.0,
        "w_0": -1.0,
        "w_a": 0.0,
        "fnl": 100.0,
        "gnl": 0.0,
        "k_max":10.0, 
        "tau_reio":0.06,
        "WhichSpectrum":"EH",
        "w0_fld":-1.0,
        "wa_fld":0.0,
    }