import os
import subprocess
from datetime import datetime


def load_template(template_path: str):
    output_template = ""
    with open(template_path, "r") as handle:
        output_template = handle.read()
    return output_template


def write_batch(batch_fpath: str, batch_content: str):
    with open(batch_fpath, "w") as handle:
        _ = handle.write(batch_content)


collectives = ["Reduce",
               "Reduce_scatter",
               "Allreduce",
               "Allgather",
               "Allgatherv",
               "Scatter",
               "Scatterv",
               "Gather",
               "Gatherv",
               "Alltoall",
               "Bcast",
               # "Barrier"
               ]

procnt = [
    18,
    36,
    54,
    72
]
mpi1_bin = "/home/hpc/ihpc/ihpc136h/workspace/prototyping/bin/IMB-MPI1"
slurm_template = load_template("templates/bench.template")

template_parameter = {"time_stamp": datetime.now().strftime("%y_%m_%d_%H-%M-%S"),
                      "job_name": "",
                      "output_dir": os.getcwd()+"/output/",
                      "err_dir": os.getcwd()+"/error/",
                      "data_dir": os.getcwd()+"/data/",
                      "n_procs": 18,
                      "off_mem_flag": "",
                      "bin": mpi1_bin
                      }

output_dir = os.getcwd()+"/output/"
err_dir = os.getcwd()+"/error/"
batch_files_dir = os.getcwd()+"/batchs/"
data_dir = os.getcwd()+"/data/"

if os.path.isdir(output_dir) == False:
    os.mkdir(output_dir)
if os.path.isdir(err_dir) == False:
    os.mkdir(err_dir)
if os.path.isdir(data_dir) == False:
    os.mkdir(data_dir)
if os.path.isdir(batch_files_dir) == False:
    os.mkdir(batch_files_dir)

log = ""

for n_procs in procnt:
    template_parameter["n_procs"] = n_procs
    for collective in collectives:
        template_parameter["job_name"] = collective
        write_batch(batch_files_dir+collective+".sh",
                    slurm_template.format(**template_parameter))
        result = subprocess.run(["sbatch", batch_files_dir+collective+".sh"],
                                capture_output=True, text=True)
        log += f"#{collective} {n_procs}" + "\n"
        log += "\tSTDOUT:" + result.stdout + "\n"
        log += "\tSTDERR:" + result.stderr + "\n"
print(log)
_ = subprocess.run(["./clean.sh"])