#!/usr/bin/env python

import sys, os, shlex
import multiprocessing as mp
import subprocess as sp
from pli.lib.util import log


def find_files(root):
    outdir = 'tmpout'
    for curdir, dirs, files in os.walk(root):
        protein_fname = None
        ligand_fname = None
        for f in files:
            if f.endswith('_protein.pdb'):
                sym = f[1:3]
                protein_fname = f'{curdir}/{f}'
            elif f.endswith('_ligand.mol2'):
                ligand_fname = f'{curdir}/{f}'
        if protein_fname and ligand_fname:
            bname = os.path.basename(protein_fname).split('_')[0]
            sym = bname[1:3]
            if not os.path.exists(f'{outdir}/{sym}/{bname}_ALP.pkl.gz'):
                yield protein_fname, ligand_fname, sym


def worker(args):
    protein_iname, ligand_iname, sym = args
    cmd = f'python pli/bin/plifinder.py {protein_iname} {ligand_iname} -o tmpout/{sym}'
    log(cmd)
    proc = sp.Popen(shlex.split(cmd), universal_newlines=True, stdout=sp.PIPE, stderr=sp.STDOUT)
    ret = proc.wait()
    return ret, protein_iname, ligand_iname, sym


def main():
    root = 'v2015'
    pool = mp.Pool(mp.cpu_count())
    count = 0
    for ret, protein_iname, ligand_iname, sym in pool.imap(worker, find_files(root)):
        if ret != 0:
            log(f'!! Error {protein_iname} {ligand_iname}')
            continue
        count += 1
        log(f'{count} {protein_iname} {ligand_iname} -> {sym}')


if __name__ == '__main__':
    main()
