from slurm_utils import create_and_submit_batch_job
from datetime import datetime
import os
from pathlib import Path
import argparse

PROJECT_ROOT = Path(os.path.realpath(__file__)).parents[1]

parser = argparse.ArgumentParser()
parser.add_argument('--interactive', action='store_true')
args = parser.parse_args()

# DATA_DIR = 'uspto_full'
SCRIPT_DIR = 'scripts'
platform = 'lumi'
project = 'project_462000833'

slurm_args = {
    'job_dir': 'jobs',
    'job_ids_file': 'job_ids.txt',
    'output_dir': 'output',
    'platform': platform,
    'project': project,
    'time': '04:00:00',
    'partition': 'small',
    'nodes': 1,
    'ntasks-per-node': 1,
    'cpus-per-task': 1,
    'mem': '100G', # 50G not enough for uspto_full
    'with_containers': True,
    'container': 'multiguide2.sif',
    'venv_path': 'multiguide2-container',
    'start_array_job': 0, # 5 to 37
    'end_array_job': 0 #37
}

time_stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
split_name_to_process = 'val'
experiment_name = f'extract_synthesis_score_{split_name_to_process}_{time_stamp}'  

script_args = {"script_dir": SCRIPT_DIR,
               "use_torchrun": 'false',
               "args": {
                   "classifier_guidance": "synthesis",
                   "classifier_guidance.dataset.split_name_to_process": split_name_to_process
               }}

script_args['script_name'] = 'extract_synthesis_scores.py'
task = f'extract_synthesis_score_{split_name_to_process}'
slurm_args['job_name'] = f'{task}_{experiment_name}'
output = create_and_submit_batch_job(slurm_args, script_args, interactive=args.interactive)