import argparse

from data.load_data import SupportedDatasets
from train.train_conf import train_conf

'''
scp 1500011335@115.27.161.31:yangshuwen/GeometryEncoding/log/qm7-conf/* log/qm7-conf
'''
parser = argparse.ArgumentParser()
parser.add_argument('--seed', type=int, default=0)
parser.add_argument('--generate', type=str, default='rdkit')
parser.add_argument('--derive', type=str, default='newton')
parser.add_argument('--compare', type=str, default='adj3')
args = parser.parse_args()
seed = args.seed
generate = args.generate
derive = args.derive
compare = args.compare

if seed:
    token = f'{generate}-{derive}-{compare}@{seed}'
    dataset_token = f'phi-psi@{seed}'
else:
    token = f'{generate}-{derive}-{compare}'
    dataset_token = 'phi-psi'

train_conf(
    special_config={
        'GENERATE_TYPE': generate,
        'DERIVE_TYPE': derive,
        'COMPARE_TYPE': compare,
        'COMPARE_MIDDLE': True,
    },
    token=token,
    dataset_token=dataset_token,
    seed=seed,
    force_save=False,
    use_tqdm=True,
    use_cuda=True
)
