#!/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('--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 = []
    for nodes in (2, 16, 64, 256):
        for channels in (1, 2):
            for rate, frames in ((44100, 31), (48000, 128), (96000, 1024)):
                result = subprocess.check_output([str(a.binary.resolve()), *map(str, (nodes, channels, rate, frames, a.blocks))], text=True)
                case = json.loads(result)
                cases.append(case)
                print(f'{nodes:3} nodes {channels}ch {rate}/{frames}: compile {case["compile_ms"]:.3f} ms, '
                      f'mean {case["mean_us"]:.3f} us, p99 {case["p99_us"]:.3f} us, load {case["load"]:.3f}')
    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.')
        key = lambda c: (c['nodes'], c['channels'], c['rate'], c['frames'])
        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()