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

start_array_job = 0
end_array_job = 0

PROJECT_ROOT = Path(os.path.realpath(__file__)).parents[1]
SCRIPT_DIR = 'scripts'
slurm_args = get_platform_info(use_gpu=False)
slurm_args.update({
    'use_srun': False,
    'job_dir': 'jobs',
    'job_ids_file': 'job_ids.txt',
    'output_dir': 'output',
    'time': '24:00:00',
    'nodes': 1,
    'ntasks-per-node': 1,
    'cpus-per-task': 1,
    'gpus-per-node': 0,
    'mem': '100G', # 50G not enough for uspto_full
    'start_array_job': start_array_job, 
    'end_array_job': end_array_job
})
classifier_guidance = 'reaction_type'
similarity_type = ''
completion_lower_limit = '0.8'
max_augmentations = '5'
data_dir = 'uspto_190/train_unique_reactions'
subset = 'val_fraction0.4_threshold10_dropped.csv'
search = 'retro_star'
classifier_dataset_name = 'uspto_190_fraction0.4_threshold10_with_product' # output_path
classifier_dataset_subset = 'train_completion0.8_augment5.csv'
single_step_model_checkpoint_full_path = os.path.join(PROJECT_ROOT, 'checkpoints', 'rsmiles_uspto_190_fraction1.0_thresh500_dropped_PtoR_aug5', 'model.product-reactants_step_250000.pt')
filter_by_vocab = 'true'

time_stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
script_args = {"script_dir": SCRIPT_DIR,
                "use_torchrun": 'false',
                "args": {
                    "search": search,
                    "classifier_guidance": classifier_guidance,
                    "classifier_guidance.similarity_type": similarity_type,
                    "classifier_guidance.dataset.completion_lower_limit": completion_lower_limit,
                    "classifier_guidance.dataset.data_dir": classifier_dataset_name,
                    "classifier_guidance.dataset.subset": classifier_dataset_subset,
                    "classifier_guidance.dataset.max_augmentations": max_augmentations,
                    #"classifier_guidance.dataset.separator": separator,
                    "reaction_dataset.data_dir": data_dir,
                    "reaction_dataset.subset": subset,
                    "reaction_dataset.filter_by_vocab": filter_by_vocab,
                    "single_step_model.checkpoint_full_path": single_step_model_checkpoint_full_path
                }
                }
script_args['script_name'] = 'get_reactant_for_reaction_type_classifier.py'
slurm_args['job_name'] = f'get_reactant_for_reaction_type_classifier_{time_stamp}'
output = create_and_submit_batch_job(slurm_args, script_args, interactive=slurm_args['interactive'])