#!/usr/bin/env python
# split input files into several temporary files

from __future__ import division
import sys
import os
import optparse
import math



o = optparse.OptionParser()
o.disable_interspersed_args()
#o.add_option("-r", "--reduce", dest="reduce",
#             default="cat $FILES")
o.add_option("-i", "--input", dest="input",
             action="append",
             help="input file")
#o.add_option("-o", "--output", dest="output",
#             help="output file")
o.add_option("-n", "--numgroup", dest="numgroup",
             type="int", default=10,
             help="number of groups to make (default: 10)")
o.add_option("-g", "--groupsize", dest="groupsize",
             type="int",
             help="size of each group")
o.add_option("-w", "--wait", dest="wait", action="store_true",
             help="wait for each command to finish")
o.add_option("-t", "--tmpprefix", dest="tmpprefix",
             help="temp output prefix")
o.add_option("-s", "--submit", dest="submit",
             default='bsub -o "$FILE.bout" -E "touch \"$FILE.bout\""  -K ',
             help="submit command")

o.add_option("", "--no-pre-clean", dest="no_pre_clean",
             action="store_true",
             help="Do not clean tempdir prior to run")
o.add_option("", "--post-cleanup", dest="post_clean",
             action="store_true",
             help="Clean up tempdir after run")
o.add_option("--clean-cmd", dest="clean_cmd",
             default='rm -f "$FILE" "$FILE.out" "$FILE.bout"',
             help="cleanup command after run")

o.add_option("-v", "--verbose", dest="verbose",
             action="store_true", default=False)


#=============================================================================

def error(text):
    sys.stderr.write(text + "\n")


def doublequote(text):
    """Wrap argument in doublequotes with propper escaping"""
    return '"' + text.replace("\\", "\\\\").replace('"', '\\"').replace("$", "\\$") + '"'


def make_cmd(cmd, **args):
    cmd2 = ""
    for key, val in args.items():
        cmd2 += 'export %s=%s; ' % (key, doublequote(val))
    return cmd2 + cmd


def main(argv):
    options, args = o.parse_args(argv[1:])

    # determine temp output prefix
    if options.tmpprefix is None:
        if "RUNPAR_PREFIX" in os.environ:
            options.tmpprefix = os.environ["RUNPAR_PREFIX"]
        else:
            error(("%s: you must specify tempdir with -t or with the " + 
                   "environment variable RUNPAR_PREFIX") % \
                    os.path.basename(sys.argv[0]))
            return 2

    # ensure temp space exists
    if not os.path.isdir(options.tmpprefix):
        # XXX: tmpprefix might contain a filename prefix
        # tmpdir = os.path.dirname(options.tmpprefix)
        tmpdir = options.tmpprefix
        if not os.path.exists(tmpdir):
            if options.verbose:
                print >>sys.stderr, "making temp directory '%s'" % tmpdir
            os.makedirs(tmpdir)
        

    # clean all previous tempfiles, if requested
    if not options.no_pre_clean:
        for fn in os.listdir(options.tmpprefix):
            fn2 = os.path.join(options.tmpprefix, fn)
            if options.verbose:
                print >>sys.stderr, "removing", fn2
            os.remove(fn2)


    # determine map command
    if len(args) > 0:
        cmd = " ".join(args)
        if cmd == "":
            error("must give map command")
            return 2
    else:
        cmd = sys.stdin.read()

    
    if not options.wait:
        options.submit = 'bsub -o "$FILE.bout" -E "touch \"$FILE.bout\""  '
        

    # determine infiles
    infiles = []
    if len(options.input) > 0:
        for arg in options.input:
            if arg == "-":
                infiles.append(sys.stdin)
            else:
                infiles.append(open(arg))
    else:
        error("must specify input file(s)")

    # determine all input lines
    lines = []
    for infile in infiles:
        lines.extend(infile.readlines())

    # determine group size
    if options.groupsize is None:
        options.groupsize = int(math.ceil(len(lines) /
                                          float(options.numgroup)))

    if len(lines) == 0:
        return 0

    # split lines
    pids = set()
    outfiles = []
    infiles = []
    for i in xrange(0, len(lines), options.groupsize):
        fn = os.path.join(options.tmpprefix, str(i))
        infiles.append(fn)
        outfiles.append(fn + ".out")

        # create input files
        out = open(fn, "w")
        for line in lines[i:i+options.groupsize]:
            out.write(line)
        out.close()

        # run command
        cmd2 = make_cmd(options.submit + doublequote(
            make_cmd('( ' + cmd + ' ) > "$FILE.out"', FILE=fn)),
                        FILE=fn)

        
        if options.verbose:
            print >>sys.stderr, cmd2

        if not options.wait:
            os.spawnlp(os.P_WAIT, "bash", "bash", "-c", cmd2)
        else:
            pids.add(os.spawnlp(os.P_NOWAIT, "bash", "bash", "-c", cmd2))

    # wait for all jobs to finish
    while len(pids) > 0:
        pid, code = os.waitpid(0, 0)
        if pid in pids:
            pids.remove(pid)
        


    # run reduce command
    '''
    outfiles = " ".join(map(quote, outfiles))
    if options.output is not None:
        out_cmd = ' > %s' % doublequote(options.output)
    else:
        out_cmd = ""
    reduce_cmd = make_cmd(options.reduce + out_cmd, FILES=outfiles)
    if options.verbose:
        print >>sys.stderr, reduce_cmd
    os.system(reduce_cmd)
    '''

    # clean up
    if options.post_clean:
        for fn in infiles:
            clean_cmd = make_cmd(options.clean_cmd, FILE=fn)
            if options.verbose:
                print >>sys.stderr, clean_cmd
            os.system(clean_cmd)

    return 0

sys.exit(main(sys.argv))
