import os
import re
import shutil
from collections import defaultdict

# Calculate number of processes
def n_processes(do_numa, n_nodes, n_slaves, scale_factor):
    import numa
    max_nodes = numa.get_max_node() + 1        
    if do_numa:
        n_nodes = max_nodes if n_nodes == -1 else n_nodes
    else:
        n_nodes = 1 if n_nodes == -1 else n_nodes

    if n_slaves != -1:
        n_sl = n_slaves
    else:
        import multiprocessing
        n_sl = int(float(multiprocessing.cpu_count()) / float(max_nodes) * scale_factor + 0.5)
    return n_nodes, n_sl

# Utility functions 
def tpe(total_events, total_time, norm_rate, events, time, n_cpu, n_workers):
    rate = total_events / total_time 
    if norm_rate == None:
        norm_rate = rate
    else:
        norm_rate /= float(n_workers)
    norm = rate / norm_rate
    
    tpe = time / events if time != 0.0 else 0.
    tpe /= norm
    return tpe * float(n_cpu) / float(n_workers)

def get_tpes(log_files, norm_rate, n_cpu, n_workers, timing_width = 60):
    average_times = defaultdict(list)

    total_events = 0
    total_time = 0
    table_start = 0
    
    logs = {}
    re_timing = re.compile(r"(?:\|?\s*(\d+(?:\.\d+)?)\s*)")
    re_name = re.compile(r"TIMER\.TIMER:(\s*)([\w0-9]+(?:\s\w+)?)\s*\|")
    re_space = re.compile(r"\s{20:}")
    if n_workers != len(log_files):
        print 'WARNING, number of log files %d is not the same as number of requested workers %d.' % (len(log_files), n_workers)
    
    for log_file in log_files:
        tmp_file = os.path.join(os.path.dirname(log_file), "tmp.log")
        with open(log_file) as input_log:
            line_number = 0
            output_log = open(tmp_file, 'w')
            for line in input_log:
                line = line[:-1]
                if ('-' * 20) in line:
                    line = 'TIMER.TIMER ' + '-' * (timing_width + 70)
                if line.startswith('TIMER'):
                    re_space.sub(" ", line)
                    try:
                        numbers = [float(n) for n in re_timing.findall(line)]
                    except TypeError:
                        print 'could not make numbers: ', re_timing.findall(line)
                        raise
                    if 'EVENT LOOP' in line:
                        total_events, total_time = numbers[-2:]
                        table_start = line_number
                    if total_events and numbers:
                        t = 1000 * tpe(total_events, total_time, norm_rate, *numbers[-2:], n_cpu = n_cpu, n_workers = n_workers)
                        line = '%s%s' % (line, '{0:>12.5f} |'.format(t))
                        try:
                            r = re_name.search(line)
                            indent = len(r.group(1))
                            element_name = r.group(2)
                            average_times[(element_name, line_number - table_start, indent)].append(t)
                        except AttributeError:
                            pass
                    loc = line.find('|')
                    if loc > timing_width:
                        line = line[: timing_width] + line[loc :]
                output_log.write(line + '\n')
                line_number += 1
            output_log.close()
        shutil.move(tmp_file, log_file)
    return average_times
