#!/usr/bin/env python
"""
    Input: Molecules in SDF, SMILES ...
    Output: Molecule file filtered with obgrep.
    Copyright 2013, Bjoern Gruening and Xavier Lucas
"""
import argparse
import multiprocessing
import os
import shlex
import shutil
import subprocess
import tempfile


def parse_command_line():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--infile', required=True, help='Molecule file.')
    parser.add_argument('-q', '--query', required=True, help='Query file, containing different SMARTS in each line.')
    parser.add_argument('-o', '--outfile', required=True, help='Path to the output file.')
    parser.add_argument("--iformat", help="Input format, like smi, sdf, inchi")
    parser.add_argument("--n-times", dest="n_times", type=int,
                        default=0, help="Print a molecule only if the pattern occurs # times inside the molecule.")
    parser.add_argument('-p', '--processors', type=int, default=multiprocessing.cpu_count())
    parser.add_argument("--invert-matches", dest="invert_matches", action="store_true",
                        default=False, help="Invert the matching, print non-matching molecules.")
    parser.add_argument("--only-name", dest="only_name", action="store_true",
                        default=False, help="Only print the name of the molecules.")
    parser.add_argument("--full-match", dest="full_match", action="store_true",
                        default=False, help="Full match, print matching-molecules only when the number of heavy atoms is also equal to the number of atoms in the SMARTS pattern.")
    parser.add_argument("--number-of-matches", dest="number_of_matches", action="store_true",
                        default=False, help="Print the number of matches.")
    return parser.parse_args()


results = list()


def mp_callback(res):
    results.append(res)


def mp_helper(query, args):
    """
        Helper function for multiprocessing.
        That function is a wrapper around obgrep.
    """

    cmd_list = []
    if args.invert_matches:
        cmd_list.append('-v')
    if args.only_name:
        cmd_list.append('-n')
    if args.full_match:
        cmd_list.append('-f')
    if args.number_of_matches:
        cmd_list.append('-c')
    if args.n_times:
        cmd_list.append('-t %s' % str(args.n_times))

    tmp = tempfile.NamedTemporaryFile(delete=False)
    cmd = 'obgrep %s "%s" %s' % (' '.join(cmd_list), query, args.infile)
    child = subprocess.Popen(shlex.split(cmd), stdout=open(tmp.name, 'w+'), stderr=subprocess.PIPE)

    stdout, stderr = child.communicate()
    return (tmp.name, query)


def obgrep(args):
    temp_file = tempfile.NamedTemporaryFile()
    temp_link = "%s.%s" % (temp_file.name, args.iformat)
    temp_file.close()
    os.symlink(args.infile, temp_link)
    args.infile = temp_link

    pool = multiprocessing.Pool(args.processors)
    for query in open(args.query):
        pool.apply_async(mp_helper, args=(query.strip(), args), callback=mp_callback)
        # mp_callback(mp_helper(query.strip(), args))
    pool.close()
    pool.join()

    out_handle = open(args.outfile, 'wb')
    for result_file, query in results:
        res_handle = open(result_file, 'rb')
        shutil.copyfileobj(res_handle, out_handle)
        res_handle.close()
        os.remove(result_file)
    out_handle.close()

    os.remove(temp_link)


def __main__():
    """
        Multiprocessing obgrep search.
    """
    args = parse_command_line()
    obgrep(args)


if __name__ == "__main__":
    __main__()