#!/usr/bin/env python3
import subprocess


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

#SBATCH --job-name=gol.out
#SBATCH --output=gol-{version}-{ntasks}.out.o%j
#SBATCH --ntasks={ntasks}
#SBATCH --cpus-per-task=1
#SBATCH --partition=shared-gpu
#SBATCH --gpus=rtx:{ntasks}
#SBATCH --mem-per-cpu={mem_per_cpu}
#SBATCH --time={time}

srun /home/users/c/coudrayb/projet-de-bachelor/game_of_life/cmake-build-release/game_of_life_{version}_bench {ntasks if ntasks < 16 else 8} 30000 30000
"""


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

#SBATCH --job-name=gol.out
#SBATCH --output=gol-{version}-{ntasks}.out.o%j
#SBATCH --ntasks={ntasks}
#SBATCH --cpus-per-task=1
#SBATCH --partition=public-bigmem
#SBATCH --mem-per-cpu={mem_per_cpu}
#SBATCH --time={time}

srun /home/users/c/coudrayb/projet-de-bachelor/game_of_life/cmake-build-release/game_of_life_{version}_bench 1 30000 30000
"""


HOURS = ["00:20: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, 3):
            ntasks = 2 ** i
            mem_per_cpu = "8000" if ntasks <= 2 else "4000"
            time = HOURS[0]
            if version == "seq" or version == "mc":
                content = slurm_cpu(version, ntasks, mem_per_cpu, time)
            else:
                content = slurm_gpu(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()