#!/usr/bin/python

import sys, subprocess, shutil
import itertools as it
from UserString import MutableString
import argparse


def chunks(l, n):
    """ Yield successive n chunks from l.
    """
    step = len(l) / n
    if step == 0:
        step = 1
    pos = 0
    itr = 0
    while pos < len(l):
        if itr < n - 1:
            yield l[pos:pos+step]
            pos += step
        else:
            yield l[pos:]
            pos = len(l)
        itr += 1


def make_wccl_cmd(corpus_paths, operator_path):
    operator_dir = shutil.os.path.dirname(operator_path)
    generic_command = "wccl-run -i poliqarp %s %s -F test-0 -V 1 -A 1 --tabs -P %s"
    corpora_string = ""
    for p in corpus_paths:
        corpora_string += "-I %s " % p

    return generic_command % (corpora_string, operator_path, operator_dir)

def load_lines(path):
    lines = []
    f = open(path, 'r')
    for w in f:
        lines.append(w.strip())
    f.close()
    return lines

def prepare_run(lexicon_words, tasks, share_dir, operator_path, ssh_config, corpora_paths):
    task_paths = []
    for chunk, task in it.izip(chunks(lexicon_words, tasks), it.count()):
        task_dir = shutil.os.path.join(share_dir, "task-%d" % task)
        shutil.os.mkdir(task_dir)
        shutil.copy(operator_path, task_dir)

        task_lexicon = open(shutil.os.path.join(task_dir, "leksykon.txt"), 'w')
        for c in chunk:
            task_lexicon.write(c + "\n")
        task_lexicon.close()

        cmd_file = open(shutil.os.path.join(task_dir, "command"), 'w')
        cmd_file.write(make_wccl_cmd(corpora_paths, shutil.os.path.join(task_dir,
          shutil.os.path.basename(operator_path))) + "\n")
        cmd_file.close()
        task_paths.append(task_dir)
    shutil.copy("helper.py", share_dir)
    shutil.copy("triple_filter.py", share_dir)
    shutil.copy(ssh_config, share_dir)
    return task_paths

def run_parallel(task_paths, ssh_config, share_dir):
    name_ssh_config = shutil.os.path.basename(ssh_config)
    parallel = subprocess.Popen(["parallel", "--sshloginfile", name_ssh_config], stdin = subprocess.PIPE,
      cwd = share_dir)
    cmd = MutableString()
    for path in task_paths:
        cmd.append(shutil.os.path.join(share_dir, "helper.py") + (" %s\n" % path))

    parallel.communicate(str(cmd))

    for path in task_paths:
        task_name = shutil.os.path.basename(path)
        log_path = shutil.os.path.join(path, "log")
        if shutil.os.path.isfile(log_path):
            with open(log_path, "r") as log:
                last_line = log.readlines()[-1].strip()
                status1, status2 = last_line.split()[1:3]
                if int(status1) != 0 or int(status2) != 0:
                    sys.stderr.write("%s failed! See %s for details.\n" % (task_name, log_path))
                else:
                    sys.stderr.write("%s finished successfully.\n" % task_name)
        else:
            sys.stderr.write("%s failed! helper.py wasn't run.\n" % task_name)

parser = argparse.ArgumentParser()
parser.add_argument('--tasks', '-t', metavar="COUNT", type=int, required=True, help='number of tasks to create')
parser.add_argument('--shared-dir', '-s', metavar="PATH", required=True, help='path to the shared directory')
parser.add_argument('--lexicon', '-l', metavar="PATH", required=True, help='path to the lexicon file')
parser.add_argument('--sshconfig', '-S', metavar="PATH", required=True, help='path to the sshconfig file')
parser.add_argument('--corpora-list', '-c', metavar="PATH", required=True, help='path to the corpora list file')
parser.add_argument('operator', metavar="operator", help='path to the wccl operator file')

args = parser.parse_args()

share_dir = args.shared_dir
lexicon = args.lexicon
parallel_ssh_config = args.sshconfig
operator_path = args.operator
tasks = args.tasks
corpora_paths = args.corpora_list

if shutil.os.listdir(share_dir) != []:
    sys.stderr.write("error: shared directory is not empty\n")
    sys.exit(-1)

task_paths = prepare_run(load_lines(lexicon), tasks, share_dir, operator_path, parallel_ssh_config,
               load_lines(corpora_paths))
run_parallel(task_paths, parallel_ssh_config, share_dir)
