#!/usr/bin/env python
import os, sys, re
import argparse
import pydim
import atexit
import shutil

def is_valid_dir(parser, arg):
    if arg == 'network':
        return arg
    if arg != None and (not arg or not os.path.exists(arg) or not os.path.isdir(arg)):
       parser.error("The file %s does not exist!" % arg)
    else:
       return arg

parser = argparse.ArgumentParser(usage = 'usage: %(prog)s file')

parser.add_argument("user_package", nargs = 1)
parser.add_argument("--numa", action = 'store_true', dest = 'numa', default = True,
                    help = "Use NUMA binding to bind instances of all tasks to a "
                    "specific CPU and memory nodes.")
parser.add_argument("--dim-dns-node", type = str, dest = "dim_dns_node", default = "localhost",
                    help = "Hostname or IP address of the DIM dns node.")
parser.add_argument("--nodes", type = int, dest = "nodes", default = -1, choices = [-1] + range(1, 10),
                    help = "How many NUMA nodes to run on, default is all.")
parser.add_argument("--log-file", type = str, dest = "log", default = 'MooreOnlineTest.log',
                    help = "Log file to use.")
parser.add_argument("--viewers", action = "store_true", dest = "viewers",
                    help = "Run the LogViewer and buffer manager monitor.")
parser.add_argument("--ncpu-scale-factor", type = float, dest = "ncpu_scale_factor", default = 1.2,
                    help = "Scale the number of available cores by this factor to obtain the number"
                    "of Moore processes to run.")
parser.add_argument("--warmup", type = int, dest = "warmup", default = 60,
                    help = "Process events for X seconds before starting to count.")
parser.add_argument("--runtime", type = int, dest = "runtime", default = 300,
                    help = "Count for X seconds.")
parser.add_argument("--slaves", type = int, dest = "slaves", default = -1,
                    help = "Number of Moore processes to run.")
parser.add_argument("--input", type = lambda x: is_valid_dir(parser, x),
                    dest = "input", default = '/localdisk/hlt1/test')
parser.add_argument("--output", type = str, dest = "output",
                    default = '/localdisk/hlt1/test_output')
parser.add_argument("--moore-output-level", type = int, dest = "output_level",
                    default = 4, choices = range(6))
parser.add_argument("--write-mdf", dest = "write_mdf", action = "store_true",
                    default = False, help = "Write MDF files, only applies if HLT1 is run")
parser.add_argument("--moore-log", dest = "moore_log", type = str,
                    default = "", help = "Write Moore tasks output to log file.")
parser.add_argument("--mode", dest = "mode", type = str, choices = ["time", "all_data"],
                    default = "time", help = "Write Moore tasks output to log file.")
parser.add_argument("-w", "--timing-width", type = int, default = 60,
                    dest = "timing_width",
                    help = "Width of the name column in the timing table")
parser.add_argument("-p", "--partition", type = str, default = "TEST",
                    dest = "partition", help = "Which partition?")

args = parser.parse_args()

os.environ["DIM_DNS_NODE"] = args.dim_dns_node
if not pydim.dis_get_dns_node():
    print "No Dim DNS node found, problem with DIM_DNS_NODE??"
    sys.exit(1)

# if the default fifo isn't there, stuff doesn't work. Create it in that case.
if not os.path.exists("/tmp/logSrv.fifo"):
    os.mkfifo("/tmp/logSrv.fifo")
        
# Write log in directory
if os.path.sep in args.moore_log:
    log_dir, f = args.moore_log.rsplit(os.path.sep, 1)
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    log_file = os.path.join(log_dir, args.log)
else:
    log_dir = '.'
    log_file = args.log

# logging
from OfflineOnline.TestUtils import *
logger = setup_loggin(__name__, log_file)

results = None
if args.write_mdf:
    od = args.output
    if not (od and os.path.exists(od) and os.path.isdir(od)):
        raise RuntimeError("Output directory %s does not exist." % od)

input_dir = args.input
if input_dir != 'network':
    files = partition(os.listdir(input_dir), mdf_mep)
    # Check that there are either mdf or mep files, not both
    if files['REST'] or (not bool(len(files['MDF'])) ^ bool(len(files['MEP']))):
        print 'The input directory must contain either only MDF or only MEP files, not a mix.'
        print 'And no other files.'
        sys.exit(-1)

    input_type = 'MDF' if files.get('MDF', []) else 'MEP'
    prefixes = set(fn[ : 4] for fn in files[input_type])
    if len(prefixes) != 1:
        print 'All files in the input directory must have filenames with the same first 4 characters.'
        sys.exit(-1)
else:
    input_type = 'MDF'
    prefixes = ['MDF']
    
# Number of workers and number of slaves
n_nodes      = args.nodes
n_slaves     = args.slaves
scale_factor = args.ncpu_scale_factor
    
mode = args.mode
# Instantiate the benchmark processes.
from MooreTests.Benchmark import MooreBenchmark
bench_args = dict(output_level = args.output_level, user_package = args.user_package[0],
                  viewers = args.viewers, n_slaves = n_slaves, input_directory = input_dir,
                  output_directory = args.output, write_mdf = args.write_mdf,
                  n_nodes = n_nodes, scale_factor = scale_factor, mode = mode,
                  input_type = input_type, log_file = args.moore_log, file_prefix = list(prefixes)[0],
                  logger = logger, numa = args.numa, partition = args.partition,
                  timing_width = args.timing_width)

av_inst = 0.
av_inst_s = 0.

# Run benchmark
with MooreBenchmark(**bench_args) as benchmark:
    benchmark.start()

    import time
    if mode == "time" and args.warmup:
        logger.info("Letting benchmark warm up for %d seconds." % args.warmup)
        time.sleep(args.warmup)

    logger.info("Starting analysis, running for %d seconds." % args.runtime)
    benchmark.run()

    if mode == "time":
        time.sleep(args.runtime)
        logger.info("Stop counting.")
        benchmark.stop()
    else:
        benchmark.wait()

    results = benchmark.results()

    inst_rates = results['Moore'][0]['inst_rates']
    # Calculate the mean instantaneous rate for normalisation
    from math import sqrt
    av_inst  = sum(inst_rates) / len(inst_rates)
    av_inst2 = sum((r ** 2 for r in inst_rates)) / len(inst_rates)
    av_inst_s = sqrt(av_inst2 - av_inst ** 2)
    logger.info('Average instantaneous rate: %8.3f +- %8.3f.\n' % (av_inst, av_inst_s))

    import shelve
    db = shelve.open(os.path.join(log_dir, 'test_results.db'))
    db['results'] = results
    db.close()
    del db

moore_log = os.path.basename(args.moore_log)
if '.' in moore_log:
    log_base = moore_log.rsplit('.', 1)[0]
else:
    log_base = moore_log

    
# Write average log file
import multiprocessing
from MooreTests.TestUtils import *
n_nodes, n_slaves = n_processes(args.numa, n_nodes, n_slaves, scale_factor)
def filter_logs(log_file):
    if not hasattr(filter_logs, 'regex'):
        filter_logs.regex = re.compile(r'^%s_(\d(\d+)).*' % log_base)
    r = filter_logs.regex.match(log_file)
    return int(r.group(2)) if r else r

log_files = [os.path.join(log_dir, lf) for lf in os.listdir(log_dir) if filter_logs(lf)]
log_files = sorted(log_files)
average_times = get_tpes(log_files, av_inst, multiprocessing.cpu_count(), n_nodes * n_slaves,
                         args.timing_width)

with open(os.path.join(log_dir, 'averages.log'), 'w') as average_log:
    average_log.write('Averate rate: %f +- %f\n' % (av_inst, av_inst_s))
    max_name_len = max(len(k[0]) for k in average_times.iterkeys())
    max_indent = max(k[2] for k in average_times.iterkeys())
    lines = sorted(average_times.iteritems(), key = lambda item: item[0][1])
    for (name, n, indent), values in lines:
        av_time  = sum(values) / len(values)
        av_time2 = sum((r ** 2 for r in values)) / len(values)
        av_time_s = sqrt(av_time2 - av_time ** 2)
        pat = "{0:<" + str(max_name_len + max_indent - indent) + "s}"
        line = (' ' * indent) + pat.format(name) + ("  %8.3f +- %5.3f\n" % (av_time, av_time_s))
        average_log.write(line)
