#!/usr/bin/env python3
"""Run the published EMA benchmark locally. No telemetry, no installation, no source upload."""
import argparse
import hashlib
import json
import os
import platform
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

SITE = 'https://taktcycles.com/showcase/talib/'
HARNESS = 'https://github.com/rdmitry0911/takt-harness/releases/download/v0.2.1/'
HARNESS_SHA256 = 'fe979337105dd50ee1d1579a3d15071c08475708eb4644324639a5244562dcbd'


def fetch(url, path):
    subprocess.run(['curl', '-fLsS', '--retry', '2', '--max-time', '120', url, '-o', str(path)], check=True)


def checked(url, path, digest):
    fetch(url, path)
    if hashlib.sha256(path.read_bytes()).hexdigest() != digest:
        raise SystemExit('Checksum mismatch: ' + path.name)
    path.chmod(0o755)


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--cpu', type=int, help='one CPU in your allowed affinity mask')
    ap.add_argument('--out', type=Path, help='new directory for the binaries and readings')
    ap.add_argument('--no-counters', action='store_true', help='measure wall time without hardware counters')
    args = ap.parse_args()
    if platform.system() != 'Linux' or platform.machine() != 'x86_64' or platform.libc_ver()[0] != 'glibc':
        raise SystemExit('The benchmark requires Linux x86-64 with glibc.')
    flags = Path('/proc/cpuinfo').read_text().split('flags', 1)[-1].splitlines()[0].split()
    if not {'avx2', 'fma'}.issubset(flags):
        raise SystemExit('The published benchmark requires AVX2 and FMA.')
    if not shutil.which('curl'):
        raise SystemExit('curl is required.')
    allowed = os.sched_getaffinity(0)
    cpu = args.cpu if args.cpu is not None else min(allowed - {8} or allowed)
    if cpu not in allowed:
        raise SystemExit('CPU is outside your allowed affinity mask.')
    os.sched_setaffinity(0, {cpu})
    if args.out:
        work = args.out.resolve()
        work.mkdir(parents=True, exist_ok=False)
    else:
        work = Path(tempfile.mkdtemp(prefix='takt-demo-'))
    print('Public EMA demo. Files and readings stay in', work, flush=True)
    fetch(SITE + 'SHA256SUMS', work / 'SHA256SUMS')
    manifest = {}
    for line in (work / 'SHA256SUMS').read_text().splitlines():
        digest, name = line.split(maxsplit=1)
        manifest[name[2:] if name.startswith('./') else name] = digest
    programs = []
    for label in ['official', 'tuned', 'takt']:
        name = 'talib-bench-' + label
        checked(SITE + 'bin/' + name, work / name, manifest['bin/' + name])
        programs.append(work / name)
    fetch(SITE + 'LICENSE', work / 'LICENSE')
    harness = work / 'takt-harness'
    checked(HARNESS + 'takt-harness-x86_64-linux', harness, HARNESS_SHA256)
    record = work / 'readings.json'
    cmd = [str(harness), 'run', '--reps', '9', '--warmup', '1', '--out', str(record)]
    if args.no_counters:
        cmd.append('--no-counters')
    for label in ['official', 'tuned', 'takt']:
        cmd += ['--label', label]
    cmd.append('--')
    for i, program in enumerate(programs):
        if i:
            cmd.append(':::')
        cmd += [str(program), 'sweep', 'EMA']
    subprocess.run(cmd, check=True)
    bundle = json.loads(record.read_text())
    for baseline in ['official', 'tuned']:
        pair = dict(bundle)
        pair['commands'] = [next(c for c in bundle['commands'] if c['label'] == label) for label in [baseline, 'takt']]
        path = work / (baseline + '-vs-takt.json')
        path.write_text(json.dumps(pair, indent=2) + '\n')
        subprocess.run([str(harness), 'compare', str(path)], check=True)
    print('\nUse the smaller speedup of the two comparisons. This is a public example, not a forecast for your code.')
    print('Keep or remove the directory when finished:', work)


if __name__ == '__main__':
    try:
        main()
    except subprocess.CalledProcessError as e:
        sys.exit(e.returncode)