import os

import sys

import glob

def pdbs_to_pdbqts(pdb_dir, pdbqt_dir, dataset):

    for file in glob.glob(os.path.join(pdb_dir, '*.pdb')):

        name = os.path.splitext(os.path.basename(file))[0]

        outfile = os.path.join(pdbqt_dir, name + '.pdbqt')

        pdb_to_pdbqt(file, outfile, dataset)

        print('Wrote converted file to {}'.format(outfile))

def pdb_to_pdbqt(pdb_file, pdbqt_file, dataset):

    if os.path.exists(pdbqt_file):

        return pdbqt_file

    if dataset == 'crossdocked':

        os.system('prepare_receptor4.py -r {} -o {}'.format(pdb_file, pdbqt_file))

    elif dataset == 'bindingmoad':

        os.system('prepare_receptor4.py -r {} -o {} -A checkhydrogens -e'.format(pdb_file, pdbqt_file))

    else:

        raise NotImplementedError

    return pdbqt_file

if __name__ == '__main__':

    pdbs_to_pdbqts(sys.argv[1], sys.argv[2], sys.argv[3])