##################################
#                                #
# Last modified 2023/02/08       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import string
import os

def run():

    if len(sys.argv) < 6:
        print('usage: python %s <commands list file> CPUs memory walltime queue prefix [-gpu] [-srun] [-long] [-batch N]') % sys.argv[0]
        print('\t The batch option will create batches of jobs on a single line separated by ;')
        sys.exit(1)

# sacct -l

    input = sys.argv[1]
    CPUs = sys.argv[2]
    mem = sys.argv[3]
    walltime = sys.argv[4]
    queue = sys.argv[5]
    prefix = sys.argv[6]

    doGPU = False
    if '-gpu' in sys.argv:
        doGPU = True

    doSR = False
    if '-srun' in sys.argv:
        doSR = True

    doLong = False
    if '-long' in sys.argv:
        doLong = True

    BS = 1
    doBatch = False
    if '-batch' in sys.argv:
        doBatch = True
        BS = int(sys.argv[sys.argv.index('-batch') + 1])

    cmd = 'wc -l ' + input
    p = os.popen(cmd, "r")
    CL = int(p.readline().split(' ')[0])

    CMDLINE = ''

    linelist = open(input)
    i=0
    for line in linelist:
        if line.strip() == '':
            continue
        i += 1
        if i % BS == 0 or i == CL:
            outfilename = prefix + '.' + str(i) + '.sbatch'
            outfile = open(outfilename,'w')
            outfile.write('#!/bin/bash' + '\n')
            outfile.write('#SBATCH --job-name=' + prefix + '-' + str(i) + '\n')
            outfile.write('#SBATCH --output=' + prefix + '.' + str(i) + '.out' '\n')
            outfile.write('#SBATCH --error=' + prefix + '.' + str(i) + '.err' '\n')
            outfile.write('#SBATCH --time=' + walltime + '\n')
            outfile.write('#SBATCH -p ' + queue + '\n')
            outfile.write('#SBATCH --cpus-per-task=' + CPUs + '\n')
            outfile.write('#SBATCH --mem=' + mem + '\n')
            if line.strip().endswith ('&'):
                CMDLINE = CMDLINE + line.strip()[0:-1]
            else:
                CMDLINE = CMDLINE + line.strip()
            if doGPU:
                outfile.write('#SBATCH --gres gpu:1' + '\n')
            if doLong:
                outfile.write('#SBATCH --qos=long' + '\n')
            if doSR:
                outfile.write('srun ' + CMDLINE + '\n')
            else:
                outfile.write(CMDLINE + '\n')
            outfile.close()
            cmd = 'sbatch ' + prefix + '.' + str(i) + '.sbatch'
            print(cmd)
            os.system(cmd)
            CMDLINE = ''
        else:
            if line.strip().endswith ('&'):
                CMDLINE = CMDLINE + line.strip()[0:-1] + '; '
            else:
                CMDLINE = CMDLINE + line.strip() + '; '
            

run()