def main_scola(parsed_args):
    from args_main import parse_arguments_main
    from low_level import print_starting_module, print_message, print_ending_module, progress_bar_from_logfile, wait_until_file_exists
    from os.path import isfile
    import subprocess
    
    print_starting_module("scola", verbose=parsed_args.verbose)
    main_dict = parse_arguments_main(parsed_args)
    paramfile=parsed_args.paramfile
    if paramfile is None:
        paramfile = main_dict["paramdir"]+"parameters_"+main_dict["simname"]+".sbmy"
    log_file = main_dict["logdir"]+main_dict["simname"]+".log"
    
    nboxes_tot = int(parsed_args.N_tiles**3)

    if parsed_args.execution == "local":
        from parameters_card import parse_arguments_card
        from tqdm import tqdm
        card_dict = parse_arguments_card(parsed_args)
        print_message("Running sCOLA in local mode.", 1, "scola", verbose=parsed_args.verbose)
        
        tile_bar = tqdm(total=nboxes_tot, desc="sCOLA", unit="box", disable=(parsed_args.verbose!=1)) # The progress bar cannot work if there are print statements in the loop
        for b in range(1,nboxes_tot+1):
            if not isfile(card_dict["OutputTilesBase"]+str(b)+".h5") or parsed_args.force:
                if isfile(log_file+f"_{b}"): # Remove the preexisting log file to allow for the progress_bar to be run normally
                    from os import remove
                    remove(log_file+f"_{b}")
                if parsed_args.verbose > 1:
                    print_message(f"Running box {b}/{nboxes_tot}.", 2, "scola", verbose=parsed_args.verbose)
                command_args = ["scola", paramfile, log_file, "-b", str(b)]
                if parsed_args.verbose < 2:
                    subprocess.Popen(command_args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) # Use Popen instead of run to be able to display the progress bar
                    progress_bar_from_logfile(log_file+f"_{b}", desc=f" Box {b}/{nboxes_tot}", verbose=parsed_args.verbose, leave=False)
                else:
                    subprocess.run(command_args)

                if parsed_args.verbose > 1:
                    print_message(f"Box {b}/{nboxes_tot} finished.", 3, "scola", verbose=parsed_args.verbose)
            else:
                if parsed_args.verbose > 1:
                    print_message(f"Box {b}/{nboxes_tot} already exists. Use -F to overwrite.", 2, "scola", verbose=parsed_args.verbose)
            tile_bar.update(1)
        tile_bar.close()

        print_message("sCOLA finished.", 1, "scola", verbose=parsed_args.verbose)
    
    elif parsed_args.execution == "slurm":
        from slurm_submission import create_slurm_script, parse_arguments_slurm
        from args_main import parse_arguments_main
        import os
        print_message("Running scola in slurm mode.", 1, "scola", verbose=parsed_args.verbose)
        slurm_dict=parse_arguments_slurm(parsed_args)
        main_dict=parse_arguments_main(parsed_args)

        if have_no_tiles(parsed_args) or parsed_args.force:
            ## Submit all boxes
            print_message("Submitting all boxes.", 2, "scola", verbose=parsed_args.verbose)
            slurm_script = slurm_dict["scripts"]+"scola_"+main_dict["simname"]+".sh"

            if not isfile(slurm_script) or parsed_args.force:
                print_message(f"SLURM script {slurm_script} does not exist (or forced). Creating it.", 2, "scola", verbose=parsed_args.verbose)
                create_slurm_script(
                    slurm_template=slurm_dict["scola_template"],
                    slurm_script=slurm_script,
                    job="scola",
                    job_config_file=paramfile,
                    job_log=log_file,
                    array=(1, nboxes_tot),
                    job_name=main_dict["simname"],
                )
                print_message(f"SLURM script written to {slurm_script}.", 3, "scola", verbose=parsed_args.verbose)
            else:
                print_message(f"SLURM script {slurm_script} exists.", 2, "scola", verbose=parsed_args.verbose)
            
            print_message(f"Submitting scola job to SLURM.", 1, "scola", verbose=parsed_args.verbose)
            command_args = ["sbatch", slurm_script]

            for b in range(1,nboxes_tot+1):
                if isfile(log_file+f"_{b}"): # Remove the preexisting log file to allow for the progress_bar to be run normally
                    os.remove(log_file+f"_{b}")

            if parsed_args.verbose < 2:
                subprocess.run(command_args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            else:
                subprocess.run(command_args)

            print_message("sCOLA job submitted.", 2, "scola", verbose=parsed_args.verbose)

        else:
            ## Submit missing boxes
            missing_tiles_arrays = get_missing_tiles_arrays(parsed_args)
            print_message(f"Submitting missing boxes: {missing_tiles_arrays}.", 2, "scola", verbose=parsed_args.verbose)

            for missing_tiles in missing_tiles_arrays:
                slurm_script = slurm_dict["scripts"]+"scola_"+main_dict["simname"]+"_"+str(missing_tiles[0])+"_"+str(missing_tiles[-1])+".sh"

                if not isfile(slurm_script):
                    print_message(f"SLURM script {slurm_script} does not exist. Creating it.", 2, "scola", verbose=parsed_args.verbose)
                    create_slurm_script(
                        slurm_template=slurm_dict["scola_template"],
                        slurm_script=slurm_script,
                        job="scola",
                        job_config_file=paramfile,
                        job_log=log_file,
                        array=missing_tiles,
                        job_name=main_dict["simname"],
                    )
                    print_message(f"SLURM script written to {slurm_script}.", 3, "scola", verbose=parsed_args.verbose)
                else:
                    print_message(f"SLURM script {slurm_script} exists.", 2, "scola", verbose=parsed_args.verbose)
            
                print_message(f"Submitting scola job to SLURM.", 1, "scola", verbose=parsed_args.verbose)
                command_args = ["sbatch", slurm_script]

                for b in range(missing_tiles[0],missing_tiles[1]+1):
                    if isfile(log_file+f"_{b}"): # Remove the preexisting log file to allow for the progress_bar to be run normally
                        os.remove(log_file+f"_{b}")

                if parsed_args.verbose < 2:
                    subprocess.run(command_args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
                else:
                    subprocess.run(command_args)

                print_message("sCOLA job submitted.", 2, "scola", verbose=parsed_args.verbose)

                os.remove(slurm_script) # Remove the script after submission (because it is specific to the missing tiles)

    else:
        raise ValueError(f"Execution mode {parsed_args.execution} not recognized.")
    
    print_ending_module("scola", verbose=parsed_args.verbose)


def main_pre_scola(parsed_args):
    """
    Pre-scola function. It calls simbelmyne to generate the LPT phi1 and phi2 potential fields.
    If they already exist, it does nothing. If all tiles exist, raises error.
    """
    from os.path import isfile
    from low_level import print_starting_module, print_message, print_ending_module
    from parameters_card import parse_arguments_card
    print_starting_module("pre-scola", verbose=parsed_args.verbose)
    card_dict = parse_arguments_card(parsed_args)

    if not isfile(card_dict["OutputLPTPotential1"]) or not isfile(card_dict["OutputLPTPotential2"]) or parsed_args.force:
        if not have_all_tiles(parsed_args):
            print_message("Running pre-scola.", 1, "pre-scola", verbose=parsed_args.verbose)
            from simbelmyne import main_simbelmyne
            main_simbelmyne(parsed_args)
        else:
            raise NotImplementedError("All tiles exists, so calling simbelmyne would generate the final output instead of the LPT potentials.")
    else:
        print_message("LPT potentials already exist. Use -F to overwrite.", 1, "pre-scola", verbose=parsed_args.verbose)
    
    print_ending_module("pre-scola", verbose=parsed_args.verbose)


def main_post_scola(parsed_args):
    """
    Post-scola function. It calls simbelmyne to generate the final output from all tiles.
    If the output already exists, it does nothing. If tiles are missing, print the missing tiles.
    """
    from os.path import isfile
    from low_level import print_starting_module, print_message, print_ending_module
    from parameters_card import parse_arguments_card
    print_starting_module("post-scola", verbose=parsed_args.verbose)
    card_dict = parse_arguments_card(parsed_args)

    output_is_required = (card_dict["WriteFinalDensity"] or card_dict["WriteFinalSnapshot"])
    if not output_is_required:
        print_message("Output is not required. Skipping post-scola.", 1, "post-scola", verbose=parsed_args.verbose)
        print_ending_module("post-scola", verbose=parsed_args.verbose)
        return

    if (card_dict["WriteFinalDensity"] and not isfile(card_dict["OutputFinalDensity"])) or (card_dict["WriteFinalSnapshot"] and not isfile(card_dict["OutputFinalSnapshot"])) or parsed_args.force:
        if have_all_tiles(parsed_args):
            print_message("Running post-scola.", 1, "post-scola", verbose=parsed_args.verbose)
            from simbelmyne import main_simbelmyne
            main_simbelmyne(parsed_args)
        else:
            missing_tiles_arrays = get_missing_tiles_arrays(parsed_args)
            count_missing_tiles = sum([m[1]-m[0]+1 for m in missing_tiles_arrays])
            nboxes_tot = int(parsed_args.N_tiles**3)
            print_message(f"All tiles are not available. Missing {count_missing_tiles} out of {nboxes_tot} ({100*count_missing_tiles/nboxes_tot:.1f}%)", 1, "post-scola", verbose=parsed_args.verbose)
            print_message(f"Missing tiles: {missing_tiles_arrays}", 2, "post-scola", verbose=parsed_args.verbose)
    else:
        print_message("Output already exists. Use -F to overwrite.", 1, "post-scola", verbose=parsed_args.verbose)
    
    print_ending_module("post-scola", verbose=parsed_args.verbose)



def have_all_tiles(parsed_args):
    from os.path import isfile
    from parameters_card import parse_arguments_card

    card_dict = parse_arguments_card(parsed_args)
    nboxes_tot = int(parsed_args.N_tiles**3)
    all_tiles = True
    for b in range(1,nboxes_tot+1):
        if not isfile(card_dict["OutputTilesBase"]+str(b)+".h5"):
            all_tiles = False
    return all_tiles


def have_no_tiles(parsed_args):
    from os.path import isfile
    from parameters_card import parse_arguments_card

    card_dict = parse_arguments_card(parsed_args)
    nboxes_tot = int(parsed_args.N_tiles**3)
    no_tiles = True
    for b in range(1,nboxes_tot+1):
        if isfile(card_dict["OutputTilesBase"]+str(b)+".h5"):
            no_tiles = False
    return no_tiles


def get_missing_tiles_arrays(parsed_args):
    from os.path import isfile
    from parameters_card import parse_arguments_card

    card_dict = parse_arguments_card(parsed_args)
    nboxes_tot = int(parsed_args.N_tiles**3)
    missing_tiles_arrays = []
    in_sequence_of_missing = False
    for b in range(1,nboxes_tot+1):
        if not isfile(card_dict["OutputTilesBase"]+str(b)+".h5"):
            if not in_sequence_of_missing:
                missing_tiles_arrays.append([b])
                in_sequence_of_missing = True
        elif in_sequence_of_missing:
            missing_tiles_arrays[-1].append(b-1)
            in_sequence_of_missing = False
    if in_sequence_of_missing:
        missing_tiles_arrays[-1].append(nboxes_tot)
    for m in missing_tiles_arrays:
        assert len(m) == 2
        assert m[0] <= m[1]
    return missing_tiles_arrays


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()

    card_dict = main_parameter_card(parsed_args)
    nboxes_tot = int(parsed_args.N_tiles**3)
    main_timestepping(parsed_args)
    main_pre_scola(parsed_args)
    if parsed_args.execution == "slurm":
        wait_until_file_exists(card_dict["OutputLPTPotential2"], verbose=parsed_args.verbose, limit=5*60)
    main_scola(parsed_args)
    if parsed_args.execution == "slurm":
        for b in range(1,nboxes_tot+1):
            wait_until_file_exists(f"{card_dict['OutputTilesBase']}{b}.h5", verbose=parsed_args.verbose, limit=5*60)
    main_post_scola(parsed_args)