from os.path import isfile
from argparse import ArgumentParser
import configparser


def register_arguments_monofonic(parser: ArgumentParser):
    """
    Register the arguments for the monofonIC parameters.
    """
    parser.add_argument("-mfc", "--monofonic_config", type=str, default=None, help="Path to the monofonIC configuration file.")
    parser.add_argument("-mfo", "--monofonic_output", type=str, default=None, help="Path to the monofonIC output file.")
    parser.add_argument("-mfset", "--monofonic_cosmo_set", type=str, default=None, help="Set of cosmological parameters to use in monofonIC.")


def parse_arguments_monofonic(parsed_args):
    """
    Parse the arguments for the monofonIC parameters.
    """
    from args_main import parse_arguments_main
    from parameters_card import parse_arguments_card_for_ICs
    from cosmo_params import parse_arguments_cosmo

    main_dict = parse_arguments_main(parsed_args)
    card_dict = parse_arguments_card_for_ICs(parsed_args)
    cosmo_dict = parse_arguments_cosmo(parsed_args)

    monofonic_dict = dict(
        config=parsed_args.monofonic_config,
        seed=main_dict["seed"],
        nthreads=main_dict["nthreads"],
        gridres=card_dict["N_LPT_mesh"],
        boxlength=card_dict["L"],
        ParameterSet=parsed_args.monofonic_cosmo_set,
        output=parsed_args.monofonic_output,
        **cosmo_dict
    )

    if monofonic_dict["config"] is None:
        monofonic_dict["config"]=main_dict["paramdir"]+"monofonic_config.conf"
    if monofonic_dict["output"] is None:
        monofonic_dict["output"]=main_dict["workdir"]+"initial_conditions_"

    return monofonic_dict



def create_monofonic_config(filename, config_params):
    """
    Creates a MUSIC configuration file based on provided parameters.
    
    :param filename: Name of the output config file.
    :param config_params: Dictionary containing configuration parameters.
    """
    config = configparser.ConfigParser(allow_no_value=True)
    config.optionxform = str  # Preserve case sensitivity
    
    for section, params in config_params.items():
        config[section] = {}
        for key, value in params.items():
            if isinstance(value, bool):
                value = 'yes' if value else 'no'
            config[section][key] = str(value)
    
    with open(filename, 'w') as configfile:
        config.write(configfile, space_around_delimiters=False)


def get_config_from_dict(monofonic_dict):
    config={}

    config["setup"] = {
        "GridRes": monofonic_dict["gridres"],
        "BoxLength": monofonic_dict["boxlength"],
        "zstart": 999.0,
        "LPTorder": 2,
        "DoBaryons": False,
        "DoBaryonVrel": False,
        "DoFixing": False,
        "DoInversion": False,
        "ParticleLoad": "sc"
    }

    ## WARNING: the cosmo dict is not updated accordingly to the ParameterSet
    if monofonic_dict["ParameterSet"] is not None:
        config["cosmology"] = {
            "ParameterSet": monofonic_dict["ParameterSet"],
            "transfer": "CLASS",
            "ztarget": 2.5,
            "ZeroRadiation": False
        }
    else:
        config["cosmology"]={
            "Omega_m": monofonic_dict["Omega_m"],
            "Omega_b": monofonic_dict["Omega_b"],
            "Omega_L": monofonic_dict["Omega_q"],
            "H0": monofonic_dict["h"]*100.,
            "n_s": monofonic_dict["n_s"],
            "sigma_8": monofonic_dict["sigma8"],
            "A_s": monofonic_dict["A_s"],
            "Tcmb": monofonic_dict["Tcmb"],
            "k_p": monofonic_dict["k_p"],
            "N_ur": monofonic_dict["N_ur"],
            "m_nu1": monofonic_dict["m_nu1"],
            "m_nu2": monofonic_dict["m_nu2"],
            "m_nu3": monofonic_dict["m_nu3"],
            "w_0": monofonic_dict["w_0"],
            "w_a": monofonic_dict["w_a"],
            "fnl": monofonic_dict["fnl"],
            "gnl": monofonic_dict["gnl"],
            "transfer": "CLASS",
            "ztarget": 2.5,
            "ZeroRadiation": False
        } 

    config["random"] = {
        "generator": "NGENIC",
        "seed": monofonic_dict["seed"]
    }

    config["execution"] = {
        "NumThreads": monofonic_dict["nthreads"]
    }

    config["output"] = {
        "format": "simbelmyne",
        "filename": monofonic_dict["output"]
    }

    return config



def main_parameters_monofonic(parsed_args):
    from low_level import print_message

    print_message("Parsing arguments for the config file.", 1, "monofonic", verbose=parsed_args.verbose)
    monofonic_dict = parse_arguments_monofonic(parsed_args)
    if isfile(monofonic_dict["config"]) and not parsed_args.force:
        print_message(f"Configuration file {monofonic_dict['config']} already exists. Use -F to overwrite.", 1, "monofonic", verbose=parsed_args.verbose)
        return monofonic_dict
    create_monofonic_config(monofonic_dict["config"], get_config_from_dict(monofonic_dict))
    print_message(f"Configuration file written to {monofonic_dict["config"]}", 2, "monofonic", verbose=parsed_args.verbose)
    
    return monofonic_dict
    

if __name__ == "__main__":
    from args_main import register_arguments_main
    from parameters_card import register_arguments_card_for_ICs
    from cosmo_params import register_arguments_cosmo

    parser = ArgumentParser(description="Create monofonIC configuration file.")
    register_arguments_main(parser)
    register_arguments_monofonic(parser)
    register_arguments_card_for_ICs(parser)
    register_arguments_cosmo(parser)
    parsed_args = parser.parse_args()

    main_parameters_monofonic(parsed_args)