#!/usr/bin/env python3
import subprocess


def slurm_gpu(file, version, ntasks, mem_per_cpu, time):
    return f"""#!/bin/sh

#SBATCH --job-name={file[0]}.out
#SBATCH --output={file[0]}-{version}-{ntasks}.out.o%j
#SBATCH --ntasks={ntasks}
#SBATCH --cpus-per-task=1
#SBATCH --partition=shared-gpu
#SBATCH --gpus=rtx:{min(ntasks, 8)}
#SBATCH --mem-per-cpu={mem_per_cpu}
#SBATCH --time={time}

srun {file[1]}_{version}_bench {min(ntasks, 8)} 30000 30000
"""


def slurm_cpu(file, version, ntasks, mem_per_cpu, time):
    return f"""#!/bin/sh

#SBATCH --job-name={file[0]}.out
#SBATCH --output={file[0]}-{version}-{ntasks}.out.o%j
#SBATCH --ntasks={ntasks}
#SBATCH --cpus-per-task=1
#SBATCH --partition=shared-cpu
#SBATCH --mem-per-cpu={mem_per_cpu}
#SBATCH --time={time}

srun {file[1]}_{version}_bench 1 30000 30000
"""


FILES = [
    ["elem"], ["/home/users/c/coudrayb/projet-de-bachelor/elementary/cmake-build-release/elementary"],
    ["gol"], ["/home/users/c/coudrayb/projet-de-bachelor/game_of_life/cmake-build-release/game_of_life"],
    ["lbm"], ["/home/users/c/coudrayb/projet-de-bachelor/lattice_boltzmann/cmake-build-release/lattice_boltzmann"]
]
HOURS = ["12:00:00", "06:00:00", "03:00:00", "02:00:00", "01:00:00", "01:00:00", "01:00:00"]


def main():
    for version in ["cuda"]:
        for i in range(0, 8):
            ntasks = 2 ** i
            mem_per_cpu = "8000" if ntasks <= 2 else "4000"
            file = FILES[0]
            time = HOURS[i]
            if version == "seq" or version == "mc":
                content = slurm_cpu(file, version, ntasks, mem_per_cpu, time)
            else:
                content = slurm_gpu(file, version, ntasks, mem_per_cpu, time)
            filename = f"srun_{version}_{ntasks}.sh"
            with open(filename, "w") as fh:
                fh.write(content)
            subprocess.run(["/usr/bin/sbatch", filename])


if __name__ == '__main__':
    main()