import sys

from openbabel import openbabel, pybel


def main():
    if len(sys.argv) == 3:
        process(sys.argv[1], sys.argv[2])
    else:
        print("Usage: convert_pdbqt_to_sdf.py <input-pdbqt-file> <output-sdf-file>")
        exit(1)


def add_property(mol, prop_name, prop_value):
    newData = openbabel.OBPairData()
    newData.SetAttribute(prop_name)
    newData.SetValue(prop_value)
    mol.OBMol.CloneData(newData)


def process(input, output):
    docked = pybel.readfile("pdbqt", input)
    sdf = pybel.Outputfile("sdf", output, overwrite=True)
    for mol in docked:
        if mol.OBMol.HasData("REMARK"):
            remark = mol.OBMol.GetData("REMARK").GetValue()
            lines = remark.splitlines()
            tokens = lines[0].split()

            # add the score property
            add_property(mol, "SCORE", tokens[2])
            # add the first RMSD property
            add_property(mol, "RMSD_LB", tokens[3])
            # add the second RMSD property
            add_property(mol, "RMSD_UB", tokens[4])

        sdf.write(mol)

    sdf.close()


if __name__ == "__main__":
    main()