from argparse import ArgumentParser

def register_arguments_main(parser:ArgumentParser):
    """
    Register the arguments for the main parameters.
    """
    parser.add_argument("-d", "--directory", type=str, default="./", help="Main directory where the output will be saved (if other dir and filenames are not specified).")
    parser.add_argument("-pd", "--paramdir", type=str, default=None, help="Directory where the parameters are saved (default is -d/params).")
    parser.add_argument("-rd", "--resultdir", type=str, default=None, help="Directory where the results are saved (default is -d/results).")
    parser.add_argument("-wd", "--workdir", type=str, default=None, help="Directory where the work is done (default is -d/work).")
    parser.add_argument("-ld", "--logdir", type=str, default=None, help="Directory where the logs are saved (default is -d/logs).")
    parser.add_argument("-n", "--simname", type=str, default=None, help="Name of the simulation (for outputs and parameter card).")
    parser.add_argument("-M", "--mode", type=str, default="tCOLA", help="Mode of the program: ICs, PM, LPT, tCOLA, sCOLA, pre_sCOLA, post_sCOLA, allPM, alltCOLA, allsCOLA, TS.")
    parser.add_argument("-ic","--ICs_gen", type=str, default="ext", help="Initial conditions generator: ext, sbmy, monofonic")
    parser.add_argument("-v", "--verbose", type=int, default=1, help="Verbosity level (0: no output, 1: basic output, 2: subprograms output).")
    parser.add_argument("--seed", type=int, default=None, help="Seed for the random number generator.")
    parser.add_argument("--nthreads", type=int, default=None, help="Number of threads to use.")
    parser.add_argument("-E","--execution", type=str, default="local", help="Execution mode: local, slurm.")
    parser.add_argument("-F","--force", action="store_true", help="Force the regeneration of all files (parameters and outputs).")

def parse_arguments_main(parsed_args):
    """
    Parse the arguments for the main parameters.
    """
    from pathlib import Path
    main_dict = dict(
            directory=parsed_args.directory,
            paramdir=parsed_args.paramdir,
            resultdir=parsed_args.resultdir,
            workdir=parsed_args.workdir,
            logdir=parsed_args.logdir,
            mode=parsed_args.mode,
            verbose=parsed_args.verbose,
            simname=parsed_args.simname,
            ICs_gen=parsed_args.ICs_gen,
            seed=parsed_args.seed,
            nthreads=parsed_args.nthreads,
            execution=parsed_args.execution,
            force=parsed_args.force,
            )
    
    ## Now for all the directories, if they are not specified, they are set using the main directory
    if main_dict["paramdir"] is None:
        main_dict["paramdir"] = main_dict["directory"]+"params/"
    if main_dict["resultdir"] is None:
        main_dict["resultdir"] = main_dict["directory"]+"results/"
    if main_dict["workdir"] is None:
        main_dict["workdir"] = main_dict["directory"]+"work/"
    if main_dict["logdir"] is None:
        main_dict["logdir"] = main_dict["directory"]+"logs/"
    if main_dict["simname"] is None:
        main_dict["simname"] = "sCOLA" if "sCOLA" in main_dict["mode"] else main_dict["mode"]
    if main_dict["seed"] is None:
        import numpy as np
        main_dict["seed"] = np.random.randint(0, 2**32)
    if main_dict["nthreads"] is None:
        import os
        main_dict["nthreads"] = os.cpu_count()

    # Create missing directories
    Path(main_dict["paramdir"]).mkdir(parents=True, exist_ok=True)
    Path(main_dict["resultdir"]).mkdir(parents=True, exist_ok=True)
    Path(main_dict["workdir"]).mkdir(parents=True, exist_ok=True)
    Path(main_dict["logdir"]).mkdir(parents=True, exist_ok=True)
    
    return main_dict