#!/bin/env python

'''Testlet to run the LatencyBandwidth test.'''

import os
import time
import test_utils as util

Processor_Numbers = [2, 4, 8, 16]

name = 'Running latency/bandwidth test'

def test(logfile):
    result = True

    (cluster, domain) = util.get_hostname()

    # get python to run
    python_env_var = os.getenv('PYTHON')

    # compile C code
    obj_file = 'ctiming_%s' % cluster
    cmd = ('mpicc ctiming.c -lm -o %s' % obj_file)
    util.log_print_nl(logfile, 'Compiling ctiming.c:')
    util.log_print_nl(logfile,cmd)
    (_, fd) = os.popen4(cmd)
    output = fd.read()
    status = fd.close()
    if status:
        util.log_print_nl(logfile, 'ERRORS COMPILING')
        util.log_print(logfile, output)
        return False

    util.log_print_nl(logfile)

    # run the tests
    results = []
    machines_file = os.path.join('~', '.machines_%s' % cluster)
    for num_procs in Processor_Numbers:
        cmd = ('mpirun -nolocal -np %2d -hostfile %s %s'
               % (num_procs, machines_file, obj_file))
        util.log_print(logfile, cmd + ' ...')
        (_, fd) = os.popen4(cmd)
        output = fd.readlines()
        status = fd.close()
        if status:
            util.log_print_nl(logfile, 'ERRORS')
            util.log_print(logfile, '\n'.join(output))
            return False

        estimated_bandwidth = None
        estimated_latency = None
        for line in output:
            if line.startswith('Estimated bandwith'):
                (_, estimated_bandwidth) = line.split(': ', 1)
                estimated_bandwidth = estimated_bandwidth.strip()
                (estimated_bandwidth, _) = estimated_bandwidth.split(' ', 1)
            if line.startswith('Estimated latency'):
                (_, estimated_latency) = line.split(': ', 1)
                estimated_latency = estimated_latency.strip()
                (estimated_latency, _) = estimated_latency.split(' ', 1)

        if estimated_bandwidth is None or estimated_latency is None:
            util.log_print_nl(logfile, 'ERRORS')
            util.log_print(logfile, '\n'.join(output))
            return False
            
        results.append((num_procs, float(estimated_bandwidth), int(estimated_latency)))
        util.log_print_nl(logfile, 'done')

    # interpret the results
    util.log_print(logfile, '\nResults of latency/bandwidth test on %s\n' % cluster)
    util.log_print(logfile, '%s\t%s\t%s\n' % ('#procs', 'b/width', 'latency'))
    
    for (num_processors, estimated_bandwidth, estimated_latency) in results:
        util.log_print(logfile, '%d\t%.1f\t%d\n'
                       % (num_processors, estimated_bandwidth, estimated_latency))

    os.remove(obj_file)

    return result


if __name__ == '__main__':
    import sys

    logfile = 'test.log'
    if len(sys.argv) > 1:
        logfile = sys.argv[1]

    try:
        os.remove(logfile)
    except:
        pass

    if not test(logfile):
        sys.exit(10)

    sys.exit(0)