import numpy as np
import sys

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]

# print("hello: ", sim_pack)
# print(task)

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 com_cap == 'gpu':
                job_file.write('#SBATCH --cluster=gmerlin6\n')
                job_file.write('#SBATCH --hint=nomultithread\n')
                job_file.write('#SBATCH --nodes=1\n')
                job_file.write('#SBATCH --ntasks=1\n')
                if sim_pack == 'qulacs':
                    job_file.write('#SBATCH --partition=gpu-short\n')
                    job_file.write('#SBATCH --gpus=geforce_rtx_2080_ti:1\n')
                else:
                    job_file.write('#SBATCH --partition=gwendolen\n')
                    job_file.write('#SBATCH --account=gwendolen\n')
                    job_file.write('#SBATCH --gpus=A100:1\n')
                if N < 30:
                    job_file.write('#SBATCH --mem=30G\n')
                else:
                    if sim_pack == 'qulacs':
                        job_file.write('#SBATCH --mem=100G\n')
                    else:
                        job_file.write('#SBATCH --mem=900G\n')

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

            elif com_cap == 'gpu1' or com_cap == 'gpu2' or com_cap == 'gpu4' or com_cap == 'gpu8':
                job_file.write('#SBATCH --cluster=gmerlin6\n')
                job_file.write('#SBATCH --hint=nomultithread\n')
                job_file.write('#SBATCH --nodes=1\n')
                ngpus = int(com_cap[-1])
                if sim_pack == 'cuquantum_qiskit':
                    job_file.write('#SBATCH --ntasks={}\n'.format(ngpus))                
                job_file.write('#SBATCH --partition=gwendolen\n')
                job_file.write('#SBATCH --account=gwendolen\n')
                ngpus = int(com_cap[-1])
                job_file.write('#SBATCH --gpus=A100:{}\n'.format(ngpus))
                if N < 30:
                    job_file.write('#SBATCH --mem=30G\n')
                else:
                    job_file.write('#SBATCH --mem=900G\n')

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

            else:

                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')

            # job_file.write('module load anaconda\n')
            # job_file.write('module load gcc/12.1.0\n')
            # job_file.write('conda activate sim_{}\n'.format(sim_pack))
            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')

            elif sim_pack == 'qulacs':
                if com_cap == 'st':
                    job_file.write('OMP_NUM_THREADS=1\n')
                    job_file.write('export OMP_NUM_THREADS\n')
                    job_file.write('QULACS_NUM_THREADS=1\n')
                    job_file.write('export QULACS_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('QULACS_NUM_THREADS=84\n')
       	       	    job_file.write('export QULACS_NUM_THREADS\n')

            elif sim_pack == 'yao':
                if com_cap == 'st':
                    job_file.write('OMP_NUM_THREADS=1\n')
                    job_file.write('export OMP_NUM_THREADS\n')
                    job_file.write('JULIA_NUM_THREADS=1\n')
                    job_file.write('export JULIA_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('JULIA_NUM_THREADS=84\n')
                    job_file.write('export JULIA_NUM_THREADS\n')


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

            if sim_pack == 'hiq':
                job_file.write('export APPTAINER_BIND=\"/data/user/gangap_a/{}_singular/{}/bash_scripts_{}_{}/:/brf,/data/user/gangap_a/{}_singular/{}/run_files_{}_{}/:/sim_lib/hiq/HiQsimulator/examples/run_files_{}_{},/data/user/gangap_a/{}_singular/{}/data_{}_{}\"\n'.format(task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec, com_cap, prec, task, sim_pack, com_cap, prec))
                job_file.write('singularity exec /data/user/gangap_a/singularity_images/hiq_qucos_sim.sif bash /brf/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, N))

            elif sim_pack == 'intel_qs_cpp':
                job_file.write('export APPTAINER_BIND=\"/data/user/gangap_a/{}_singular/{}/bash_scripts_{}_{}/:/brf,/data/user/gangap_a/{}_singular/intel_qs_cpp/:/home\"\n'.format(task, sim_pack, com_cap, prec, task))
                # if N==N_arr[0]:
                    # job_file.write('singularity exec /data/user/gangap_a/singularity_images/qucos_sim.sif bash /brf/{}_{}_{}.sh\n'.format(task, com_cap, prec))
                    # job_file.write('sleep 100\n')
                job_file.write('singularity exec /data/user/gangap_a/singularity_images/qucos_sim.sif bash /brf/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, N))
                # else:
                #     job_file.write('singularity exec /data/user/gangap_a/singularity_images/qucos_sim.sif bash /brf/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, N))
                
                
            elif sim_pack == 'quest':
                # job_file.write('export APPTAINER_BIND=\"/data/user/gangap_a/{}_singular/{}/bash_scripts_{}_{}/:/brf,/data/user/gangap_a/{}_singular/{}/run_files_{}_{}/:/sim_lib/quest/QuEST/{}/run_files_{}_{}\"\n'.format(task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec, task, com_cap, prec))
                job_file.write('export APPTAINER_BIND=\"/data/user/gangap_a/{}_singular/{}/bash_scripts_{}_{}/:/brf,/data/user/gangap_a/{}_singular/{}/run_files_{}_{}/:/home,/data/user/gangap_a/{}_singular/quest/build_{}_{}:/sim_lib/quest/QuEST/build_{}_{}\"\n'.format(task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec, task, com_cap, prec, com_cap, prec))
                if com_cap == 'st' or com_cap == 'mt':
                    job_file.write('singularity exec /data/user/gangap_a/singularity_images/qucos_sim.sif bash /brf/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, N))
                elif 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))

            elif sim_pack == 'qrack_sch' or sim_pack == 'qrack_opt':
                if com_cap == 'st' or com_cap == 'mt':
                    job_file.write('export APPTAINER_BIND=\"/data/user/gangap_a/{}_singular/{}/bash_scripts_{}_{}/:/brf,/data/user/gangap_a/{}_singular/{}/run_files_{}_{}/:/home,/data/user/gangap_a/{}_singular/{}/bin_st_{}:/bin_st_{}\"\n'.format(task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec, task, sim_pack, prec, prec))
                    job_file.write('singularity exec /data/user/gangap_a/singularity_images/qucos_sim.sif bash /brf/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, N))
                elif com_cap == 'gpu':
                    job_file.write('export APPTAINER_BIND=\"/data/user/gangap_a/{}_singular/{}/bash_scripts_{}_{}/:/brf,/data/user/gangap_a/{}_singular/{}/run_files_{}_{}/:/home,/data/user/gangap_a/{}_singular/{}/bin_gpu_{}:/bin_gpu_{}\"\n'.format(task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec, task, sim_pack, prec, prec))
                    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)) 


            elif sim_pack == 'qpp':
                if com_cap == 'st' or com_cap == 'mt':
                    job_file.write('export APPTAINER_BIND=\"/data/user/gangap_a/{}_singular/{}/bash_scripts_{}_{}/:/brf,/data/user/gangap_a/{}_singular/{}/run_files_{}_{}/:/home,/data/user/gangap_a/{}_singular/{}/bin_st_{}:/bin_st_{}\"\n'.format(task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec, task, sim_pack, prec, prec))
                    job_file.write('singularity exec /data/user/gangap_a/singularity_images/qucos_sim.sif bash /brf/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, N))

            # elif sim_pack == 'pennylane_l' and com_cap == 'gpu':
            elif sim_pack == 'qulacs' and com_cap == 'gpu':
                 # job_file.write('module load cuda/11.5.1\n')
                 job_file.write('export APPTAINER_BIND=\"/data/user/gangap_a/{}_singular/{}/bash_scripts_{}_{}/:/brf,/data/user/gangap_a/{}_singular/{}/run_files_{}_{}/:/home,/data/user/gangap_a/{}_singular/{}/data_{}_{}/"\n'.format(task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec))
                 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))

            elif sim_pack == 'yao' and com_cap == 'gpu':
                 job_file.write('module load cuda/11.5.1\n')
                 job_file.write('bash /data/user/gangap_a/{}_singular/yao/bash_scripts_{}_{}/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, task, com_cap, prec, N))

            elif sim_pack == 'pennylane_l' and com_cap == 'gpu':
                 job_file.write('module load anaconda\n')
                 job_file.write('conda activate sim_pennylane_gpu\n')
                 job_file.write('module load cuda/11.5.1\n')
                 job_file.write('bash /data/user/gangap_a/{}_singular/pennylane_l/bash_scripts_{}_{}/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, task, com_cap, prec, N))

            elif sim_pack == 'qsimcirq' and com_cap == 'gpu':
                 job_file.write('module load anaconda\n')
                 job_file.write('conda activate sim_qsimcirq_a100\n')
                 job_file.write('module load cuda/11.4.3\n')
                 job_file.write('bash /data/user/gangap_a/{}_singular/qsimcirq/bash_scripts_{}_{}/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, task, com_cap, prec, N))

            elif sim_pack == 'qpanda' and com_cap == 'gpu':
                 job_file.write('module load anaconda\n')
                 job_file.write('conda activate sim_qpanda_gpu\n')
                 job_file.write('module load cuda/11.4.3\n')
                 job_file.write('bash /data/user/gangap_a/{}_singular/qpanda/bash_scripts_{}_{}/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, task, com_cap, prec, N))

            elif sim_pack == 'cuquantum_qiskit' and (com_cap == 'gpu' or com_cap == 'gpu1' or com_cap == 'gpu2' or com_cap == 'gpu4' or com_cap == 'gpu8'):
                 job_file.write('export APPTAINER_BIND=\"/data/user/gangap_a/{}_singular/{}/bash_scripts_{}_{}/:/brf,/data/user/gangap_a/{}_singular/{}/run_files_{}_{}/:/home,/data/user/gangap_a/{}_singular/{}/data_{}_{}/\"\n'.format(task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec))                 
                 job_file.write('singularity exec --nv /data/user/gangap_a/singularity_images/sim_cuq_nvidia.simg bash /brf/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, N))

            elif sim_pack == 'cuquantum_qsimcirq' and (com_cap == 'gpu' or com_cap == 'gpu1' or com_cap == 'gpu2' or com_cap =='gpu4' or com_cap == 'gpu8'):
                 job_file.write('export APPTAINER_BIND=\"/data/user/gangap_a/{}_singular/{}/bash_scripts_{}_{}/:/brf,/data/user/gangap_a/{}_singular/{}/run_files_{}_{}/:/home,/data/user/gangap_a/{}_singular/{}/data_{}_{}/\"\n'.format(task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec))                 
                 job_file.write('singularity exec --nv /data/user/gangap_a/singularity_images/sim_cuq_nvidia.simg bash /brf/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, N))

            elif sim_pack == 'cuda_quantum': # or com_cap == 'gpu1' or com_cap == 'gpu2' or com_cap =='gpu4' or com_cap == 'gpu8'):
                 if com_cap == 'st':
                     job_file.write('OMP_NUM_THREADS=1\n')
                     job_file.write('export OMP_NUM_THREADS\n')
                     job_file.write('export APPTAINER_BIND=\"/data/user/gangap_a/{}_singular/{}/bash_scripts_{}_{}/:/brf,/data/user/gangap_a/{}_singular/{}/run_files_{}_{}/:/home,/data/user/gangap_a/{}_singular/{}/data_{}_{}/\"\n'.format(task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec))                 
                     job_file.write('singularity exec /data/user/gangap_a/singularity_images/sim_cudaq_nvidia.sif bash /brf/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, 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 APPTAINER_BIND=\"/data/user/gangap_a/{}_singular/{}/bash_scripts_{}_{}/:/brf,/data/user/gangap_a/{}_singular/{}/run_files_{}_{}/:/home,/data/user/gangap_a/{}_singular/{}/data_{}_{}/\"\n'.format(task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec))                 
                     job_file.write('singularity exec /data/user/gangap_a/singularity_images/sim_cudaq_nvidia.sif bash /brf/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, N))
                 elif com_cap == 'gpu':
                     job_file.write('export APPTAINER_BIND=\"/data/user/gangap_a/{}_singular/{}/bash_scripts_{}_{}/:/brf,/data/user/gangap_a/{}_singular/{}/run_files_{}_{}/:/home,/data/user/gangap_a/{}_singular/{}/data_{}_{}/\"\n'.format(task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec))                 
                     job_file.write('singularity exec --nv /data/user/gangap_a/singularity_images/sim_cudaq_nvidia.sif bash /brf/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, N))

            elif sim_pack == 'qibojit' and (com_cap == 'gpu1' or com_cap == 'gpu2' or com_cap =='gpu4' or com_cap == 'gpu8'):
                 job_file.write('export APPTAINER_BIND=\"/data/user/gangap_a/{}_singular/{}/bash_scripts_{}_{}/:/brf,/data/user/gangap_a/{}_singular/{}/run_files_{}_{}/:/home,/data/user/gangap_a/{}_singular/{}/data_{}_{}/\"\n'.format(task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec))
                 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:
                
                # if com_cap == 'gpu':
                #     job_file.write('module load cuda/11.5.1\n')
                job_file.write('export APPTAINER_BIND=\"/data/user/gangap_a/{}_singular/{}/bash_scripts_{}_{}/:/brf,/data/user/gangap_a/{}_singular/{}/run_files_{}_{}/:/home,/data/user/gangap_a/{}_singular/{}/data_{}_{}/\"\n'.format(task, sim_pack, com_cap, prec, task, sim_pack, com_cap, prec, task, sim_pack, 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 /data/user/gangap_a/singularity_images/qucos_sim.sif bash /brf/{}_{}_{}_{}.sh\n'.format(task, com_cap, prec, N))