from sh import prun

class DataCollector(object):
    def __init__(self):
        self.data = 0.0

    def __call__(self, line):
        if 'Took' in line:
            self.data = float(line.split()[1])

class FakeCollector(object):
    def __init__(self, s):
        self.data = s

def run_single(nodes, procs, i, t, out):
    return prun('-np', nodes, '-%d' % (procs,), '-sge-script', '/cm/shared/package/reserve.slurm/etc/prun-openmpi',
            'assign2_1', i, t, _out=out, _bg=True)

def run_suite(nodes, procs):
    collectors = [FakeCollector("Data for nodes=%d procs=%d" % (nodes, procs))]
    processes = []
    for i in range(0, 6):
        dc = DataCollector()
        collectors.append(dc)
        p = run_single(nodes, procs, 10**(3+i), 2*(10**(7-i)), out=dc)
        processes.append(p)

    return collectors, processes

def wait_for_suites():
    [p.wait() for p in processes]

    for collector in collectors:
        print collector.data

collectors, processes = [], []

def add_suite(nodes, procs):
    c, p = run_suite(nodes, procs)
    collectors.extend(c)
    processes.extend(p)

def run_all():
    print("Starting run...")

    for i in range(1, 9):
        add_suite(i, 8)
    for i in range(1, 9):
        add_suite(1, i)
    for i in range(1, 9):
        add_suite(i, 1)

    wait_for_suites()


if __name__ == '__main__':
    run_all()