import argparse
import requests
import sys

URL_single = "https://biosig.lab.uq.edu.au/csm_ab/api/prediction_single"
URL_pose = "https://biosig.lab.uq.edu.au/csm_ab/api/prediction_pose"


def main(args):
    job_type = args.job_type
    job_id = args.job_id

    params = {
        "job_id": job_id,
    }

    if job_type == 'single':
        req = requests.get(URL_single, json=params)
        print(req.json())
        return True
    else:
        req = requests.get(URL_pose, data=params)
        print(req.json())
        return True


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='Retrieve job results for CSM-AB')
    parser.add_argument('job_type', type=str, choices=['single', 'pose'],
                        help='choose between single and pose prediction')
    parser.add_argument('job_id', type=str,
                        help='Job identifier code generated upon submission')

    args = parser.parse_args()
    main(args)

import argparse
import requests
import sys

URL_single = "https://biosig.lab.uq.edu.au/csm_ab/api/prediction_single"
URL_pose = "https://biosig.lab.uq.edu.au/csm_ab/api/prediction_pose"

def main(args):
    if args.job_type == 'single':
        if args.pdb_accession:
            data = {"pdb_accession": args.pdb_accession}
            req = requests.post(URL_single, json=data)
            print(req.json())
        else:
            files = {"pdb_file": open(args.pdb_file.name, 'rb')}
            req = requests.post(URL_single, files=files)
            print(req.json())
    else:
        files = {'receptor_file': open(args.receptor_file.name, 'rb'), 'pose_file': open(args.pose_file.name, 'rb')}
        req = requests.post(URL_pose, files=files)
        print(req.json())


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='Job submission for CSM-AB')
    parser.add_argument('job_type', type=str, choices=['single', 'pose'], help='choose between single and pose prediction')
    parser.add_argument('--pdb_file', dest='pdb_file', type=argparse.FileType('r'), help='PDB file')
    parser.add_argument('--pdb_accession', dest='pdb_accession', type=str, help='Type PDB-ID')
    parser.add_argument('--receptor_file', dest='receptor_file', type=argparse.FileType('r'), help='PDB file')
    parser.add_argument('--pose_file', dest='pose_file', type=argparse.FileType('r'), help='PDB file')

    args = parser.parse_args()

    if args.job_type == 'single':
        if not (args.pdb_file or args.pdb_accession):
            parser.error('For single job type, either --pdb_file or --pdb_accession must be provided.')
    else:
        if not (args.receptor_file and args.pose_file):
            parser.error('For pose job type, both --receptor_file and --pose_file must be provided.')

    main(args)