import argparse

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

'''
scp 1500011335@115.27.161.31:yangshuwen/GeometryEncoding/log/qm9-conf/* log/qm9-conf
scp 1500011335@115.27.161.31:yangshuwen/GeometryEncoding/train/conf/qm9/* train/conf/qm9
scp data/dataset_pickles/qm9-phi-psi.pickle 1500011335@115.27.161.31:yangshuwen/GeometryEncoding/data/dataset_pickles
'''
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='equiv-trunc')
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'proc@{seed}'
else:
    token = f'{generate}-{derive}-{compare}'
    dataset_token = 'proc'

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