def main(parsed_args):
    from low_level import print_starting_module, print_message, print_ending_module, wait_until_file_exists
    from os.path import isfile
    from args_main import parse_arguments_main

    print_starting_module("control", verbose=parsed_args.verbose)
    main_dict = parse_arguments_main(parsed_args)

    match main_dict["mode"]:

        case "ICs" | "InitialConditions" | "InitialConditionsGenerator" | "ICsGenerator" | "ICsGen" | "ini":
            print_message("Running initial conditions generator.", 1, "control", verbose=parsed_args.verbose)
            from ICs import main_ICs
            main_ICs(parsed_args)
            print_message("Initial conditions generator finished.", 1, "control", verbose=parsed_args.verbose)
        

        case "TS" | "timestepping":
            print_message("Running timestepping generator.", 1, "control", verbose=parsed_args.verbose)
            from timestepping import main_timestepping
            main_timestepping(parsed_args)
            print_message("Timestepping generator finished.", 1, "control", verbose=parsed_args.verbose)
        

        case "PM" | "LPT" | "tCOLA" | "simbelmyne" | "sbmy":
            print_message(f"Running Simbelmyne in mode {main_dict["mode"]}.", 1, "control", verbose=parsed_args.verbose)
            from simbelmyne import main_simbelmyne
            from parameters_card import parse_arguments_card
            from os.path import isfile

            card_dict = parse_arguments_card(parsed_args)

            check_consistency(card_dict, main_dict["mode"])

            output_is_required = (card_dict["WriteFinalDensity"] or card_dict["WriteFinalSnapshot"])
            if not output_is_required:
                print_message("Output is not required. Skipping simbelmyne.", 1, "control", verbose=parsed_args.verbose)
            else:
                if (card_dict["WriteFinalDensity"] and not isfile(card_dict["OutputFinalDensity"])) or (card_dict["WriteFinalSnapshot"] and not isfile(card_dict["OutputFinalSnapshot"])) or parsed_args.force:
                    main_simbelmyne(parsed_args)
                    print_message("Simbelmyne finished.", 1, "control", verbose=parsed_args.verbose)
                else:
                    print_message("Output files already exist. Use -F to overwrite.", 1, "control", verbose=parsed_args.verbose)
        

        case "pre_sCOLA":
            print_message("Running pre-sCOLA.", 1, "control", verbose=parsed_args.verbose)
            from scola import main_pre_scola
            from parameters_card import parse_arguments_card

            card_dict = parse_arguments_card(parsed_args)

            check_consistency(card_dict, main_dict["mode"])
            
            main_pre_scola(parsed_args)
            print_message("Pre-sCOLA finished.", 1, "control", verbose=parsed_args.verbose)
        

        case "post_sCOLA":
            print_message("Running post-sCOLA.", 1, "control", verbose=parsed_args.verbose)
            from scola import main_post_scola
            from parameters_card import parse_arguments_card

            card_dict = parse_arguments_card(parsed_args)

            check_consistency(card_dict, main_dict["mode"])
            
            main_post_scola(parsed_args)
            print_message("Post-sCOLA finished.", 1, "control", verbose=parsed_args.verbose)


        case "sCOLA":        
            print_message("Running sCOLA.", 1, "control", verbose=parsed_args.verbose)
            from scola import main_scola
            from parameters_card import parse_arguments_card

            card_dict = parse_arguments_card(parsed_args)

            check_consistency(card_dict, main_dict["mode"])
            
            main_scola(parsed_args)
            print_message("sCOLA finished.", 1, "control", verbose=parsed_args.verbose)
        

        case "alltCOLA" | "allPM":
            print_message(f"Running ICs and Simbelmyne in mode {main_dict["mode"]}.", 1, "control", verbose=parsed_args.verbose)
            from parameters_card import parse_arguments_card, main_parameter_card
            from timestepping import main_timestepping
            from ICs import main_ICs
            from simbelmyne import main_simbelmyne
            from os.path import isfile
            from low_level import wait_until_file_exists

            card_dict = main_parameter_card(parsed_args)

            check_consistency(card_dict, main_dict["mode"])

            main_timestepping(parsed_args)

            ## Check consistency of ICs_gen and ICs
            if main_dict["ICs_gen"] == "monofonic":
                from parameters_monofonic import parse_arguments_monofonic
                monofonic_dict = parse_arguments_monofonic(parsed_args)
                if monofonic_dict["output"]+"DM_delta.h5" != card_dict["ICs"]:
                    raise ValueError(f"ICs {card_dict['ICs']} does not match monofonic output {monofonic_dict['output']+'DM_delta.h5'}")

            print_message("Running initial conditions generator.", 1, "control", verbose=parsed_args.verbose)
            main_ICs(parsed_args)

            ## Wait for the initial conditions to be generated
            if main_dict["execution"] == "slurm" and card_dict["ICsMode"] == 2:
                wait_until_file_exists(card_dict["ICs"])
            if main_dict["execution"] == "slurm" and card_dict["ICsMode"] == 1:
                wait_until_file_exists(card_dict["InputWhiteNoise"])
                wait_until_file_exists(card_dict["InputPowerSpectrum"])

            print_message("Initial conditions generator finished.", 1, "control", verbose=parsed_args.verbose)

            output_is_required = (card_dict["WriteFinalDensity"] or card_dict["WriteFinalSnapshot"])
            if not output_is_required:
                print_message("Output is not required. Skipping simbelmyne.", 1, "control", verbose=parsed_args.verbose)
            else:
                if (card_dict["WriteFinalDensity"] and not isfile(card_dict["OutputFinalDensity"])) or (card_dict["WriteFinalSnapshot"] and not isfile(card_dict["OutputFinalSnapshot"])) or parsed_args.force:
                    main_simbelmyne(parsed_args)
                    print_message("Simbelmyne finished.", 1, "control", verbose=parsed_args.verbose)
                else:
                    print_message("Output files already exist. Use -F to overwrite.", 1, "control", verbose=parsed_args.verbose)

            print_message(f"Running {main_dict["mode"]} finished.", 1, "control", verbose=parsed_args.verbose)
        

        case "allsCOLA":
            print_message(f"Running ICs, pre_sCOLA, sCOLA and post_sCOLA.", 1, "control", verbose=parsed_args.verbose)
            from parameters_card import parse_arguments_card, main_parameter_card
            from timestepping import main_timestepping
            from ICs import main_ICs
            from scola import main_scola, main_pre_scola, main_post_scola

            card_dict = main_parameter_card(parsed_args)

            check_consistency(card_dict, main_dict["mode"])

            main_timestepping(parsed_args)

            ## Check consistency of ICs_gen and OutputLPTPotential
            if main_dict["ICs_gen"] == "monofonic":
                from parameters_monofonic import parse_arguments_monofonic
                monofonic_dict = parse_arguments_monofonic(parsed_args)
                if monofonic_dict["output"]+"DM_phi.h5" != card_dict["OutputLPTPotential1"]:
                    raise ValueError(f"OutputLPTPotential1 {card_dict['OutputLPTPotential1']} does not match monofonic output {monofonic_dict['output']+'DM_phi.h5'}")
                if monofonic_dict["output"]+"DM_phi2.h5" != card_dict["OutputLPTPotential2"]:
                    raise ValueError(f"OutputLPTPotential2 {card_dict['OutputLPTPotential2']} does not match monofonic output {monofonic_dict['output']+'DM_phi2.h5'}")

            print_message("Running initial conditions generator.", 1, "control", verbose=parsed_args.verbose)
            main_ICs(parsed_args)

            ## Wait for the initial conditions to be generated
            if main_dict["execution"] == "slurm" and card_dict["ICsMode"] == 2:
                wait_until_file_exists(card_dict["ICs"])
            if main_dict["execution"] == "slurm" and card_dict["ICsMode"] == 1:
                wait_until_file_exists(card_dict["InputWhiteNoise"])
                wait_until_file_exists(card_dict["InputPowerSpectrum"])

            print_message("Initial conditions generator finished.", 1, "control", verbose=parsed_args.verbose)

            ## If we don't rely on monofonic to generate the initial LPT potentials, use pre_sCOLA
            if not main_dict["ICs_gen"] == "monofonic":
                print_message("Running pre-sCOLA.", 1, "control", verbose=parsed_args.verbose)
                main_pre_scola(parsed_args)
                if parsed_args.execution == "slurm":
                    wait_until_file_exists(card_dict["OutputLPTPotential2"], verbose=parsed_args.verbose, limit=5*60)
                print_message("Pre-sCOLA finished.", 1, "control", verbose=parsed_args.verbose)
            else:
                print_message("ICs_gen is monofonic. Skipping pre-sCOLA.", 1, "control", verbose=parsed_args.verbose)
            
            print_message("Running sCOLA.", 1, "control", verbose=parsed_args.verbose)
            main_scola(parsed_args)
            if parsed_args.execution == "slurm":
                from tqdm import tqdm
                from low_level import progress_bar_from_logfile
                for b in tqdm(range(1,parsed_args.N_tiles**3+1), desc="sCOLA", unit="box", disable=(parsed_args.verbose==0)):
                    progress_bar_from_logfile(main_dict["logdir"]+main_dict["simname"]+".log_"+str(b), desc=f"Box {b}/{parsed_args.N_tiles**3}", verbose=parsed_args.verbose, leave=False)
            print_message("sCOLA finished.", 1, "control", verbose=parsed_args.verbose)

            print_message("Running post-sCOLA.", 1, "control", verbose=parsed_args.verbose)
            main_post_scola(parsed_args)
            print_message("Post-sCOLA finished.", 1, "control", verbose=parsed_args.verbose)

            print_message("Running allsCOLA finished.", 1, "control", verbose=parsed_args.verbose)
            
        case _:
            raise ValueError(f"Unknown mode: {main_dict['mode']}")

    print_ending_module("control", verbose=parsed_args.verbose)




def check_consistency(card_dict, mode):
    ## Check consistency of EvolutionMode and ModulePMCOLA
    if mode == "PM" or mode == "allPM":
        if card_dict["EvolutionMode"] != 1:
            raise ValueError(f"EvolutionMode is not 1: EvolutionMode={card_dict["EvolutionMode"]}")
        if card_dict["ModulePMCOLA"] != 1:
            raise ValueError(f"ModulePMCOLA is not 1: ModulePMCOLA={card_dict["ModulePMCOLA"]}")
    elif mode == "tCOLA" or mode == "alltCOLA":
        if card_dict["EvolutionMode"] != 2:
            raise ValueError(f"EvolutionMode is not 2: EvolutionMode={card_dict["EvolutionMode"]}")
        if card_dict["ModulePMCOLA"] != 1:
            raise ValueError(f"ModulePMCOLA is not 1: ModulePMCOLA={card_dict["ModulePMCOLA"]}")
    elif mode == "LPT":
        if card_dict["ModulePMCOLA"] !=0:
            raise ValueError(f"ModulePMCOLA is not 0: ModulePMCOLA={card_dict["ModulePMCOLA"]}")
    elif mode == "pre_sCOLA" or mode == "post_sCOLA" or mode == "sCOLA" or mode=="allsCOLA":
        if card_dict["EvolutionMode"] != 3:
            raise ValueError(f"EvolutionMode is not 3: EvolutionMode={card_dict["EvolutionMode"]}")
        if card_dict["ModulePMCOLA"] != 1:
            raise ValueError(f"ModulePMCOLA is not 1: ModulePMCOLA={card_dict["ModulePMCOLA"]}")


if __name__ == "__main__":
    from argparse import ArgumentParser
    from args_main import register_arguments_main
    from timestepping import register_arguments_timestepping, main_timestepping
    from parameters_card import register_arguments_card, main_parameter_card
    from cosmo_params import register_arguments_cosmo
    from parameters_monofonic import register_arguments_monofonic
    from slurm_submission import register_arguments_slurm
    from low_level import wait_until_file_exists

    parser = ArgumentParser(description="Run sCOLA.")
    register_arguments_main(parser)
    register_arguments_timestepping(parser)
    register_arguments_monofonic(parser)
    register_arguments_slurm(parser)
    register_arguments_card(parser)
    register_arguments_cosmo(parser)
    parsed_args = parser.parse_args()
    main(parsed_args)