import pandas as pd
import numpy as np
from rdchiral import template_extractor
from tqdm import tqdm
import multiprocessing

cores = multiprocessing.cpu_count()
pool = multiprocessing.Pool(cores)

template_train = pd.read_csv('/notebooks/Codes/brs/USPTO/schneider50k/template_train.csv')
template_test = pd.read_csv('/notebooks/Codes/brs/USPTO/schneider50k/template_test.csv')
template_val = pd.read_csv('/notebooks/Codes/brs/USPTO/schneider50k/template_val.csv')

raw_data = pd.concat((template_train,template_test,template_val))



def extract(reaction):
    try:
        return template_extractor.extract_from_reaction(reaction)
    except KeyboardInterrupt:
        raise KeyboardInterrupt
    except Exception as e:
        return 'Err'

reactions = [{'_id': reaction[0], 'reactants': reaction[2].split('>>')[0], 'products': reaction[2].split('>>')[1]} for reaction in raw_data.values]

# templates = pool.map(extract, tqdm(reactions))

templates = [extract(reaction) for reaction in tqdm(reactions)]

print(123)