'''
Run this when first downloading the repo/installing the dependencies 
to make sure things are working.
'''
import pickle
import os
from multiguide.desp.DESP import DESP
from multiguide.desp.retro_predictor import RetroPredictor
from multiguide.helpers import PROJECT_ROOT

print(PROJECT_ROOT)
DEVICE = 'cpu'
RETROMODEL = os.path.join(PROJECT_ROOT, 'data', 'desp_data', 'model_retro.pt')
RETROTEMPLATES = os.path.join(PROJECT_ROOT, 'data', 'desp_data', 'idx2template_retro.json')
OUTPUT_PATH = os.path.join(PROJECT_ROOT, 'experiments', 'check_setup')
os.makedirs(OUTPUT_PATH, exist_ok=True)

retro_predictor = RetroPredictor(model_path=RETROMODEL,
                                 templates_path=RETROTEMPLATES)
# NOTE: IMPORTANT: if you're running this locally for a quick check,
# make sure the strategy is set to 'retro', since f2e and f2f require access to gpu
## specifically a piece of code in faiss throws a segmentation fault without much details
# if you have access to gpu you can test the other strategies there.
desp = DESP(strategy='retro',
            retro_predictor=retro_predictor) # switch to 'f2f' if you want to try F2F
result = desp.search(
    'COC(=O)CC12CCC(c3ccc(Br)cc3)(CC1)CO2', # Target SMILES
    ['COC(=O)C1(c2ccc(Br)cc2)CCC(=O)CC1'],  # List of starting materials SMILES
    stop_on_first_solution=False
)
print(result)

pickle.dump(result[1], open(os.path.join(OUTPUT_PATH, 'routes.pkl'), 'wb'))
pickle.dump(result[2], open(os.path.join(OUTPUT_PATH, 'searcher.pkl'), 'wb'))
desp.visualize_route(result[1][0], os.path.join(OUTPUT_PATH, 'route'))
