import argparse
import time
from parsers.parser import Parser
from parsers.config import get_config, print_cfg
from trainer import Trainer
from sampler import Sampler, Sampler_mol


def main(work_type_args):
    args = Parser().parse()
    config = get_config(args.config, args.seed)
    if args.print_cfg:
        print_cfg(config)

    # -------- Train --------
    if work_type_args.type == 'train':
        trainer = Trainer(config) 
        config.ckpt = trainer.train(time.strftime('%b%d-%H:%M:%S', time.gmtime()))
        if 'sample' in config.keys():
            assert isinstance(config.sampler.snr, list) and isinstance(config.sampler.scale_eps, list)
            config.sampler.snr.insert(0, None)
            sampler = Sampler_mol_search(config) if config.data.data in ['QM9', 'ZINC250k'] \
                        else Sampler_search(config) 
            sampler.sample()

    # -------- Generation --------
    elif work_type_args.type == 'sample':
        sampler = Sampler_mol(config) if config.data.data in ['QM9', 'ZINC250k'] else Sampler(config) 
        sampler.sample()

    else:
        raise ValueError(f'Wrong type : {work_type_args.type}')

if __name__ == '__main__':

    work_type_parser = argparse.ArgumentParser()
    work_type_parser.add_argument('--type', type=str, required=True)
    main(work_type_parser.parse_known_args()[0])
