#!/usr/bin/env python3
import sys
import os
import re
import csv

# Accept PMU events with or without user-only suffix (":u")
PMU_EVENTS = {
    'cycles': 'cycles',
    'cycles:u': 'cycles',
    'instructions': 'instructions',
    'instructions:u': 'instructions',
    'L1-dcache-load-misses': 'l1_miss',
    'L1-dcache-load-misses:u': 'l1_miss',
    'branch-misses': 'br_miss',
    'branch-misses:u': 'br_miss',
}

USDT_EVENTS = {
    'sdt:hakmem:sll_pop': 'sll_pop',
    'sdt:hakmem:mag_pop': 'mag_pop',
    'sdt:hakmem:front_pop': 'front_pop',
    'sdt:hakmem:bump_hit': 'bump_hit',
    'sdt:hakmem:slow_alloc': 'slow_alloc',
    'sdt:hakmem:sll_push': 'sll_push',
    'sdt:hakmem:mag_push': 'mag_push',
    'sdt:hakmem:spill_super': 'spill_super',
    'sdt:hakmem:spill_tiny': 'spill_tiny',
    'sdt:hakmem:remote_drain': 'remote_drain',
    'sdt:hakmem:superslab_alloc': 'super_alloc',
    'sdt:hakmem:superslab_fail': 'super_fail',
    'sdt:hakmem:quick_pop': 'quick_pop',
    'sdt:hakmem:quick_refill_sll': 'quick_refill_sll',
    'sdt:hakmem:quick_refill_mag': 'quick_refill_mag',
    'sdt:hakmem:bitmap_burst': 'bitmap_burst',
    'sdt:hakmem:mag_refill': 'mag_refill',
    'sdt:hakmem:bitmap_scan': 'bitmap_scan',
}

def parse_value(s):
    s = s.strip()
    # perf may print numbers with no separators in -x , mode; best-effort
    try:
        return int(s)
    except ValueError:
        # try float to int
        try:
            return int(float(s))
        except Exception:
            return None

def parse_stat_file(path):
    data = {}
    with open(path, 'r', errors='ignore') as f:
        for line in f:
            parts = [p.strip() for p in line.strip().split(',')]
            if len(parts) < 3:
                continue
            val = parse_value(parts[0])
            event = parts[2]
            if val is None:
                continue
            # Normalize PMU event key (strip optional ":u")
            if not event.startswith('sdt:'):
                base = event.split(':')[0]
                if event not in PMU_EVENTS and base in PMU_EVENTS:
                    event = base
            if event in PMU_EVENTS:
                data[PMU_EVENTS[event]] = val
            elif event in USDT_EVENTS:
                name = USDT_EVENTS[event]
                data[name] = data.get(name, 0) + val
            # else ignore
    return data

def main():
    if len(sys.argv) != 2:
        print("Usage: parse_usdt_stat.py <usdt_results_dir>")
        sys.exit(1)
    root = sys.argv[1]
    rows = []
    for fn in sorted(os.listdir(root)):
        if not fn.endswith('.stat.csv'):
            continue
        m = re.match(r'(?P<alloc>hakmem|system)_s(?P<size>\d+)_b(?P<batch>\d+)_c(?P<cycles>\d+)\.stat\.csv', fn)
        if not m:
            continue
        meta = m.groupdict()
        path = os.path.join(root, fn)
        stats = parse_stat_file(path)
        row = {
            'allocator': meta['alloc'],
            'size': int(meta['size']),
            'batch': int(meta['batch']),
            'cycles_param': int(meta['cycles']),
        }
        row.update(stats)
        # derived
        total_pops = sum(row.get(k, 0) for k in ('sll_pop','mag_pop','front_pop'))
        if total_pops > 0:
            row['front_rate'] = row.get('front_pop',0)/total_pops
            row['sll_rate'] = row.get('sll_pop',0)/total_pops
            row['mag_rate'] = row.get('mag_pop',0)/total_pops
        else:
            row['front_rate'] = row['sll_rate'] = row['mag_rate'] = 0.0
        rows.append(row)

    # sort for readability
    rows.sort(key=lambda r: (r['allocator'], r['size'], r['batch']))
    out = os.path.join(root, 'summary.csv')
    # collect headers
    headers = ['allocator','size','batch','cycles_param'] + list(PMU_EVENTS.values()) + list(USDT_EVENTS.values()) + ['front_rate','sll_rate','mag_rate']
    # remove duplicates but keep order
    seen = set()
    hdr_final = []
    for h in headers:
        if h not in seen:
            hdr_final.append(h)
            seen.add(h)
    with open(out, 'w', newline='') as f:
        w = csv.DictWriter(f, fieldnames=hdr_final)
        w.writeheader()
        for r in rows:
            w.writerow(r)
    print(out)

if __name__ == '__main__':
    main()