#!/usr/bin/env python
"""Script for generating slurm scripts"""

import os

from mace import Mace


TMPL_MPI = """#!/bin/bash
#SBATCH --partition=<partition>
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=<ncore>
#SBATCH --job-name=test
#SBATCH --time=12:00:00
#SBATCH --output=slurm-%j.out
#SBATCH --error=slurm-%j.err

# Load prerequisites
module load compiler mpi
"""


TMPL_OMP = """#!/bin/bash
#SBATCH --partition=<partition>
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=<ncore>
#SBATCH --job-name=default
#SBATCH --time=12:00:00
#SBATCH --output=slurm-%j.out
#SBATCH --error=slurm-%j.err

# Load prerequisites
module load compiler mkl mpi
"""


TMPL_GPU = """#!/bin/bash
#SBATCH --partition=<partition>
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=16
#SBATCH --mem=80G
#SBATCH --gres gpu:<ngpu>
#SBATCH --job-name=test
#SBATCH --time=12:00:00
#SBATCH --output=slurm-%j.out
#SBATCH --error=slurm-%j.err

# Load prerequisites
module purge
module add nvhpc/23.5
"""

TMPL_TEST = """#!/bin/bash

# Set environment variables to mimic SLURM
export SLURM_CPUS_PER_TASK=10
export SLURM_SUBMIT_DIR=$(pwd)
"""


TMPL_COMMON = """
# Set number of threads
export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK
export MKL_NUM_THREADS=$SLURM_CPUS_PER_TASK

# Change to working directory and run the job
cd $SLURM_SUBMIT_DIR
topdir=$(pwd)

# Directories of programs and scripts
tbplas_dir=$HOME/proj/tbplas-cpp
exe=$tbplas_dir/build/bin/speedtest
script=$tbplas_dir/samples/speedtest/speedtest.py

# Tasks
tasks_tbpm="haydock dos ac_cond dyn_pol dc_cond hall_cond qe wft qe_speed wft_speed"
tasks_diag="diag_all feast"

run_tasks () {
    for t in $*; do
        workdir=$topdir/$t
        test -d $workdir && rm -rf $workdir
        mkdir $workdir && cd $workdir
        mpirun -genv FI_PROVIDER=mlx $exe $t
        #$exe $t
        #python $script $t
        wait
    done
}

run_tasks $tasks_tbpm
#run_tasks $tasks_diag
"""


def write_script(file_name: str, **kwargs: dict) -> None:
    m = Mace(kwargs)
    if file_name.find("mpi") != -1:
        template = TMPL_MPI + TMPL_COMMON
    elif file_name.find("omp") != -1:
        template = TMPL_OMP + TMPL_COMMON
    elif file_name.find("gpu") != -1:
        template = TMPL_GPU + TMPL_COMMON
    else:
        template = TMPL_TEST + TMPL_COMMON
    with open(file_name, "w") as out_file:
        out_file.writelines(m.expand_lines([template]))


def main() -> None:
    write_script("mpi_c1.slurm", partition="cpu_2d", ncore=32)
    write_script("mpi_c2.slurm", partition="c2_2d", ncore=64)
    write_script("omp_c1.slurm", partition="cpu_2d", ncore=32)
    write_script("omp_c2.slurm", partition="c2_2d", ncore=64)
    write_script("gpu.slurm", partition="gpu_2d", ngpu=1)
    write_script("run_tests.sh")
    os.system("chmod +x run_tests.sh")


if __name__ == "__main__":
    main()