import numpy as np
import sys
import os

N_arr = np.arange(6, 36, 2)
# N_arr = np.append(N_arr, np.arange(41, 51, 1))
# _partition = 'hourly'

# sim_pack = 'qiskit'
# task = 'hdyn'
# com_cap = 'st'
# prec = 'dp'

sim_pack = sys.argv[1]
task = sys.argv[2]
com_cap = sys.argv[3]
prec = sys.argv[4]

gen_path = os.path.abspath(os.getcwd()) + '/'
par_dir_path = os.path.abspath(os.path.join(gen_path, os.pardir)) + '/'
sim_path = par_dir_path + '/{}/'.format(sim_pack)

for ind, N in enumerate(N_arr):
    with open('{}_{}_{}_{}.sh'.format(task, com_cap, prec, N), 'a') as job_file:
            job_file.write('#!/bin/bash -l\n')
            job_file.write('#SBATCH --job-name={}_{}_{}_{}_{}\n'.format(sim_pack, task, com_cap, prec, N))

            if N < 26:
                job_file.write('#SBATCH --partition=hourly\n')
            else:
                job_file.write('#SBATCH --partition=general\n')

            if com_cap == 'st':
                job_file.write('#SBATCH --hint=nomultithread\n')
                job_file.write('#SBATCH --nodes=1\n')
                job_file.write('#SBATCH --ntasks=1\n')
                job_file.write('#SBATCH --cpus-per-task=1\n')
                job_file.write('#SBATCH --ntasks-per-core=1\n')

            elif com_cap == 'mt':
                job_file.write('#SBATCH --hint=multithread\n')
                job_file.write('#SBATCH --nodes=1\n')
                job_file.write('#SBATCH --ntasks=1\n')
                job_file.write('#SBATCH --cpus-per-task=42\n')
                job_file.write('#SBATCH --ntasks-per-core=2\n')

            if N < 26:
                job_file.write('#SBATCH --mem=30G\n')
            else:
                job_file.write('#SBATCH --mem=300G\n')
            if N < 26:
                job_file.write('#SBATCH --time=01:00:00\n')
            else:
                job_file.write('#SBATCH --time=23:00:00\n')

            # job_file.write('#SBATCH --mail-user=amit.gangapuram@psi.ch\n')
            # job_file.write('#SBATCH --mail-type=BEGIN,END,FAIL\n')
            # job_file.write('#SBATCH --output /data/user/gangap_a/{}_singular/{}/output_{}_{}/{}_{}_{}_{}.out\n'.format(task, sim_pack, com_cap, prec, task, com_cap, prec, N))
            # job_file.write('#SBATCH --error /data/user/gangap_a/{}_singular/{}/error_{}_{}/{}_{}_{}_{}.err\n'.format(task, sim_pack, com_cap, prec, task, com_cap, prec, N))
            job_file.write('\n\n')

            if sim_pack == 'projectq' or sim_pack == 'quest' or sim_pack == 'intel_qs_cpp' or sim_pack == 'svsim' or sim_pack == 'hybridq' or sim_pack == 'hiq' or sim_pack == 'qibo' or sim_pack == 'qibojit' or sim_pack == 'qrack_sch' or sim_pack == 'qrack_opt' or sim_pack == 'pennylane_l' or sim_pack == 'qpp' or sim_pack == 'qpanda' or sim_pack =='braket' or sim_pack == 'myqlm' or sim_pack == 'myqlm_cpp':
                if com_cap == 'st':
                    job_file.write('OMP_NUM_THREADS=1\n')
                    job_file.write('export OMP_NUM_THREADS\n')

                elif com_cap == 'mt':
                    job_file.write('OMP_NUM_THREADS=84\n')
                    job_file.write('export OMP_NUM_THREADS\n')

                job_file.write('export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK\n')


            else:
                job_file.write('export APPTAINER_BIND="{}bash_scripts_{}_{}/:/brf,{}run_files_{}_{}/:/home,{}data_{}_{}/\"\n'.format(par_dir_path, com_cap, prec, par_dir_path, com_cap, prec, par_dir_path, com_cap, prec))

                if com_cap == 'gpu':
                    job_file.write('singularity exec --nv /data/user/gangap_a/singularity_images/qucos_gpu_sim.sif bash /brf/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, N))
                else:
                    job_file.write('singularity exec <path_to_singularity_image>/qucos_sim.sif bash /brf/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, N))