#!/usr/bin/env python
from __future__ import print_function

import subprocess
import re
import sys
from pprint import pprint
from itertools import chain
from copy import copy
from textwrap import dedent

import yaml

from prefixedunit import PrefixedUnit


def get_match_or_break(regex, haystack, flags=re.MULTILINE):
    m = re.search(regex, haystack, flags)
    assert m, "could not find "+repr(regex)
    return m.groups()


def get_machine_topology():
    topo = subprocess.Popen(['likwid-topology'], stdout=subprocess.PIPE).communicate()[0]
    cpuinfo = open('/proc/cpuinfo', 'r').read()
    machine = {
        'model type': get_match_or_break(r'^CPU type:\s+(.+?)\s*$', topo)[0],
        'model name': get_match_or_break(r'^model name\s+:\s+(.+?)\s*$', cpuinfo)[0],
        'sockets': int(get_match_or_break(r'^Sockets:\s+([0-9]+)\s*$', topo)[0]),
        'cores per socket': int(get_match_or_break(r'^Cores per socket:\s+([0-9]+)\s*$', topo)[0]),
        'threads per core': int(get_match_or_break(r'^Threads per core:\s+([0-9]+)\s*$', topo)[0]),
        'clock': 'INFORMATION_REQUIRED',
        'FLOPs per cycle': 'INFORMATION_REQUIRED',
        'micro-architecture': 'INFORMATION_REQUIRED',
        'icc architecture flags': 'INFORMATION_REQUIRED',
        'cacheline size': 'INFORMATION_REQUIRED',
        'overlapping ports': 'INFORAMTION_REQUIRED',
        'non-overlapping ports': 'INFORMATION_REQUIRED',
    }

    threads_start = topo.find('HWThread')
    threads_end = topo.find('Cache Topology')
    threads = {}
    for line in topo[threads_start:threads_end].split('\n'):
        m = re.match(r'([0-9]+)\s+([0-9]+)\s+([0-9]+)\s+([0-9]+)', line)
        if m:
            threads[m.groups()[0]] = (m.groups()[1:])

    cache_start = topo.find('Cache Topology')
    cache_end = topo.find('NUMA Topology')
    machine['memory hierarchy'] = []
    mem_level = {}
    for line in topo[cache_start:cache_end].split('\n'):
        if line.startswith('Level:'):
            mem_level['level'] = 'L'+line.split(':')[1].strip()
        elif line.startswith('Size:'):
            mem_level['size per group'] = PrefixedUnit(line.split(':')[1].strip())
        elif line.startswith('Cache groups:'):
            mem_level['groups'] = line.count('(')
            mem_level['cores per group'] = \
                (machine['cores per socket'] * machine['sockets']) / mem_level['groups']
            mem_level['threads per group'] = \
                mem_level['cores per group'] * machine['threads per core']
        mem_level['cycles per cacheline transfer'] = 'INFORMATION_REQUIRED'
        mem_level['bandwidth'] = 'INFORMATION_REQUIRED'

        if len(mem_level) == 8:
            machine['memory hierarchy'].append(mem_level)
            mem_level = {}
    machine['memory hierarchy'].append({
        'level': 'MEM',
        'size per group': None,
        'cores per group': machine['cores per socket'],
        'threads per group': machine['threads per core'] * machine['cores per socket'],
        'cycles per cacheline transfer': None,
        'bandwidth': 'INFORMATION_REQUIRED'
    })

    return machine


def measure_bw(type_, total_size, threads_per_core, max_threads_per_core, cores_per_socket,
               sockets, iterations=1000):
    """*size* is given in kilo bytes"""
    groups = []
    for s in range(sockets):
        groups += [
            '-w',
            'S' + str(s) + ':' + str(total_size) + 'kB:' +
            str(threads_per_core * cores_per_socket) +
            ':1:'+str(int(max_threads_per_core/threads_per_core))]
    # for older likwid versions add ['-g', str(sockets), '-i', str(iterations)] to cmd
    cmd = ['likwid-bench', '-t', type_]+groups
    output = subprocess.Popen(cmd, stdout=subprocess.PIPE).communicate()[0]
    bw = float(get_match_or_break(r'^MByte/s:\s+([0-9]+(?:\.[0-9]+)?)\s*$', output)[0])
    return PrefixedUnit(bw, 'MB/s')


def cli():
    # TODO support everything described here
    if '-h' in sys.argv[1:] or '--help' in sys.argv[1:]:
        print('''Usage:', sys.argv[0], '[-h] {collect|measure} [machinefile] | upgrade machinefile
        
        collect will retriev as much hardware information as possible, without benchmarking
        measure will do the same as collect, but also include memory benchmarks
        
        If machinefile already exists the CPU name will be compared. If they matche, measurements 
        will proceed and the file is updated accordingly. All other information in the file 
        (typically manually inserted) will be left alone.
        
        If no machinefile is given, the information will be printed to stdout.
        
        updgrade will transform machinefile to the most up-to-date machine file version.
        ''')


def main():
    machine = get_machine_topology()

    total_threads = machine['threads per core'] * machine['cores per socket']
    benchmarks = {'kernels': {}, 'measurements': {}}
    machine['benchmarks'] = benchmarks
    benchmarks['kernels'] = {
        'load': {
            'read streams': {'streams': 1, 'bytes': PrefixedUnit(8, 'B')},
            'read+write streams': {'streams': 0, 'bytes': PrefixedUnit(0, 'B')},
            'write streams': {'streams': 0, 'bytes': PrefixedUnit(0, 'B')},
            'FLOPs per iteration': 0},
        'copy': {
            'read streams': {'streams': 1, 'bytes': PrefixedUnit(8, 'B')},
            'read+write streams': {'streams': 0, 'bytes': PrefixedUnit(0, 'B')},
            'write streams': {'streams': 1, 'bytes': PrefixedUnit(8, 'B')},
            'FLOPs per iteration': 0},
        'update': {
            'read streams': {'streams': 1, 'bytes': PrefixedUnit(8, 'B')},
            'read+write streams': {'streams': 1, 'bytes': PrefixedUnit(8, 'B')},
            'write streams': {'streams': 1, 'bytes': PrefixedUnit(8, 'B')},
            'FLOPs per iteration': 0},
        'triad': {
            'read streams': {'streams': 3, 'bytes': PrefixedUnit(24, 'B')},
            'read+write streams': {'streams': 0, 'bytes': PrefixedUnit(0, 'B')},
            'write streams': {'streams': 1, 'bytes': PrefixedUnit(8, 'B')},
            'FLOPs per iteration': 2},
        'daxpy': {
            'read streams': {'streams': 2, 'bytes': PrefixedUnit(16, 'B')},
            'read+write streams': {'streams': 1, 'bytes': PrefixedUnit(8, 'B')},
            'write streams': {'streams': 1, 'bytes': PrefixedUnit(8, 'B')},
            'FLOPs per iteration': 2}, }

    USAGE_FACTOR = 0.5

    cores = range(1, machine['cores per socket']+1)
    for mem in machine['memory hierarchy']:
        measurement = {}
        machine['benchmarks']['measurements'][mem['level']] = measurement

        for threads_per_core in range(1, machine['threads per core']+1):
            threads = [c*threads_per_core for c in cores]
            if mem['size per group'] is not None:
                total_sizes = [
                    max(mem['size per group']*c/mem['cores per group'],
                        mem['size per group'])*USAGE_FACTOR
                    for c in cores]
            else:
                last_mem = machine['memory hierarchy'][-2]
                total_sizes = [last_mem['size per group']/USAGE_FACTOR for c in cores]
            sizes_per_core = [t/cores[i] for i, t in enumerate(total_sizes)]
            sizes_per_thread = [t/threads[i] for i, t in enumerate(total_sizes)]

            measurement[threads_per_core] = {
                'threads per core': threads_per_core,
                'cores': copy(cores),
                'threads': threads,
                'size per core': sizes_per_core,
                'size per thread': sizes_per_thread,
                'total size': total_sizes,
                'results': {}, }

    print('Progress: ', end='', file=sys.stderr)
    sys.stderr.flush()
    for mem_level in machine['benchmarks']['measurements'].keys():
        for threads_per_core in machine['benchmarks']['measurements'][mem_level].keys():
            measurement = machine['benchmarks']['measurements'][mem_level][threads_per_core]
            measurement['results'] = {}
            for kernel in machine['benchmarks']['kernels'].keys():
                measurement['results'][kernel] = []
                for i, total_size in enumerate(measurement['total size']):
                    measurement['results'][kernel].append(measure_bw(
                        kernel,
                        int(float(total_size)/1000),
                        threads_per_core,
                        machine['threads per core'],
                        measurement['cores'][i],
                        1,  # Sockets
                        iterations=1000))

                    print('.', end='', file=sys.stderr)
                    sys.stderr.flush()

    print(yaml.dump(machine))

if __name__ == '__main__':
    main()
