#!/usr/bin/env python3
"""Record graph timing matrix; optionally compare a matching host baseline."""
import argparse
import json
from pathlib import Path
import platform
import shlex
import subprocess
import sys


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--workload', choices=('graph', 'project'), default='graph')
    p.add_argument('--binary', type=Path, required=True)
    p.add_argument('--output', type=Path, required=True)
    p.add_argument('--baseline', type=Path)
    p.add_argument('--max-regression', type=float, default=1.5)
    p.add_argument('--noise-us', type=float, default=1.0)
    p.add_argument('--blocks', type=int, default=300)
    p.add_argument('--compiler', default='cc')
    p.add_argument('--flags', default='-O2 -g')
    a = p.parse_args()
    if not 10 <= a.blocks <= 100000 or not 1 <= a.max_regression <= 100 or not 0 <= a.noise_us <= 100:
        p.error('blocks must be 10..100000 and regression factor 1..100')
    cpu = next((line.split(':', 1)[1].strip() for line in Path('/proc/cpuinfo').read_text().splitlines()
                if line.startswith('model name')), platform.machine())
    compiler = subprocess.check_output(shlex.split(a.compiler) + ['--version'], text=True).splitlines()[0]
    context = dict(cpu=cpu, architecture=platform.machine(), compiler=compiler, flags=a.flags,
                   kernel=platform.release(), scheduler='ordinary process; no realtime scheduling')
    cases = []
    if a.workload == 'graph':
        matrix = [(nodes, channels, rate, frames, a.blocks)
                  for nodes in (2, 16, 64, 256) for channels in (1, 2)
                  for rate, frames in ((44100, 31), (48000, 128), (96000, 1024))]
    else:
        matrix = [(a.blocks, tracks, frames) for tracks in (1, 16, 64)
                  for frames in (31, 128, 1024)]
    for arguments in matrix:
        result = subprocess.check_output([str(a.binary.resolve()), *map(str, arguments)], text=True, timeout=120)
        case = json.loads(result)
        cases.append(case)
        print(f'{a.workload} {arguments}: mean {case["mean_us"]:.3f} us, '
              f'p99 {case["p99_us"]:.3f} us, misses {case["deadline_misses"]}, '
              f'CPU {case["process_cpu_ms"]:.3f} ms, peak RSS {case["peak_rss_kib"]} KiB')
    context['workload'] = a.workload
    report = dict(schema=1, context=context, cases=cases)
    a.output.parent.mkdir(parents=True, exist_ok=True)
    a.output.write_text(json.dumps(report, indent=2) + '\n')
    if a.baseline:
        baseline = json.loads(a.baseline.read_text())
        if baseline.get('schema') != 1 or any(baseline['context'][k] != context[k] for k in ('cpu', 'architecture', 'compiler', 'flags')):
            sys.exit('Baseline host/compiler/flags mismatch; record a baseline for this configuration.')
        if baseline['context'].get('workload', 'graph') != a.workload:
            sys.exit('Baseline workload mismatch.')
        fields = ('nodes', 'channels', 'rate', 'frames') if a.workload == 'graph' else ('tracks', 'notes', 'points', 'rate', 'frames')
        key = lambda c: tuple(c[field] for field in fields)
        old = {key(c): c for c in baseline['cases']}
        failed = []
        for c in cases:
            if key(c) not in old:
                failed.append(f'missing baseline case {key(c)}')
            elif c['mean_us'] > max(old[key(c)]['mean_us'] * a.max_regression, old[key(c)]['mean_us'] + a.noise_us):
                failed.append(f'{key(c)}: mean {c["mean_us"]:.3f} us exceeds {a.max_regression}x baseline')
        if failed:
            sys.exit('\n'.join(failed))
        print('Mean processing times satisfy the requested regression budget.')


if __name__ == '__main__':
    main()