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')
parser.add_argument('--platform', type=str)
args = parser.parse_args()
# DATA_DIR = 'uspto_full'
SCRIPT_DIR = 'scripts'
platform = args.platform

if platform == 'puhti':
    project = 'project_2007775'
    partition = 'small'
    with_containers = False
    puhti_module = 'pytorch/2.4'
    venv_path = '/projappl/project_2007775/rxn-insight'
    container = None
else:
    raise ValueError(f'Platform {platform} not supported')

slurm_args = {
    'job_dir': 'jobs',
    'job_ids_file': 'job_ids.txt',
    'output_dir': 'output',
    'platform': platform,
    'project': project,
    'time': '00:30:00',
    'partition': partition,
    'nodes': 1,
    'ntasks-per-node': 1,
    'cpus-per-task': 1,
    'mem': '20G', # 50G not enough for uspto_full
    'with_containers': with_containers,
    'container': container,
    'venv_path': venv_path,
    'puhti_module': puhti_module,
    'start_array_job': 0, # 5 to 37
    'end_array_job': 0 #37
}
time_stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
files = [f'route_{i}' for i in range(190)]
#files = ['all_reactions.csv']

for i, file in enumerate(files):
    script_args = {"script_dir": SCRIPT_DIR,
                    "use_torchrun": 'false',
                    "args": {
                            "reaction_dataset.data_dir": 'uspto_190/reactions',
                            "reaction_dataset.subset": file
                    }}
    script_args['script_name'] = 'apply_rxn_insight.py'
    job_name = f"apply_rxn_insight_{time_stamp}_{i}"
    slurm_args['job_name'] = job_name
    output = create_and_submit_batch_job(slurm_args, script_args, interactive=args.interactive)
# /scratch/project_462000833/multiguide/data/predictors_old