#!/usr/bin/env python3
"""Compare OpenMP schedules in a temporary, instrumented baseline renderer."""
import argparse
import gzip
import hashlib
import json
import os
from pathlib import Path
import resource
import shlex
import subprocess
import time

HERE = Path(__file__).resolve().parent
ROOT = HERE.parent.parent
BASE = '8cf6106127fc19732cee123e9890d401a6838a51'
p = argparse.ArgumentParser()
p.add_argument('--work-dir', type=Path, default=Path('/tmp/gr-ray-schedules'))
p.add_argument('--schedules', nargs='+', default=['static', 'static,32', 'dynamic,32', 'dynamic,128'])
p.add_argument('--tag', default='sweep')
p.add_argument('--production', action='store_true',
               help='Use uninstrumented baseline static or current dynamic,32 source (one schedule).')
a = p.parse_args()
a.work_dir.mkdir(parents=True, exist_ok=True)
os.chdir(ROOT)
commands = []

def source(path):
    return subprocess.check_output(['git', 'show', f'{BASE}:{path}'], text=True)

def build(cmd):
    commands.append(shlex.join(cmd))
    subprocess.run(cmd, check=True, stdout=build_log, stderr=subprocess.STDOUT)

# Only temporary copies use schedule(runtime) and diagnostic counters.
ray = source('src/ray.c').replace('#include <stdlib.h>', '#include <stdlib.h>\n#include <stdio.h>')
ray = ray.replace('#pragma omp parallel for schedule(static)\n  for (size_t i = 0; i < p->count; ++i) {', '''  const int capacity = omp_get_max_threads();
  size_t counts[capacity], step_counts[capacity];
  double work_seconds[capacity];
  for (int w = 0; w < capacity; ++w) {
    counts[w] = step_counts[w] = 0; work_seconds[w] = 0;
  }
  const double batch_start = omp_get_wtime();
#pragma omp parallel
  {
  const int worker = omp_get_thread_num();
  const double work_start = omp_get_wtime();
#pragma omp for schedule(runtime) nowait
  for (size_t i = 0; i < p->count; ++i) {''')
ray = ray.replace('    GeodesicRayState s = {.coordinate_time', '    ++counts[worker];\n    const unsigned int old_steps = p->steps[i];\n    GeodesicRayState s = {.coordinate_time')
ray = ray.replace('    p->t[i] = s.coordinate_time;', '    step_counts[worker] += s.steps - old_steps;\n    p->t[i] = s.coordinate_time;')
ray = ray.replace('\n}\n\nint ray_pool_has_live', '''
  work_seconds[worker] = omp_get_wtime() - work_start;
  }
  fprintf(stderr, "BENCH slab=%.9g pool=%zu seconds=%.9f active/steps/seconds:",
          slab->t_hi, p->count, omp_get_wtime() - batch_start);
  for (int w = 0; w < capacity; ++w)
    fprintf(stderr, " %zu/%zu/%.9f", counts[w], step_counts[w], work_seconds[w]);
  fputc('\\n', stderr);
}

int ray_pool_has_live''')
main = source('src/main.c').replace('    const int traced = trace_movie_generation(', '    const double generation_start = omp_get_wtime();\n    const int traced = trace_movie_generation(')
main = main.replace('    if (traced < 0)\n      goto done;', '    fprintf(stderr, "BENCH generation=%zu seconds=%.9f result=%d\\n", generation, omp_get_wtime() - generation_start, traced);\n    if (traced < 0)\n      goto done;')
if a.production:
    if len(a.schedules) != 1 or a.schedules[0] not in ('static', 'dynamic,32'):
        p.error('--production requires exactly one of static or dynamic,32')
    ray = (source('src/ray.c') if a.schedules[0] == 'static'
           else (ROOT/'src/ray.c').read_text())
    assert ('schedule(static)' if a.schedules[0] == 'static'
            else 'schedule(dynamic, 32)') in ray
    main = source('src/main.c')
(a.work_dir/'ray.c').write_text(ray)
(a.work_dir/'main.c').write_text(main)
with (HERE/(a.tag+'-build.log')).open('w') as build_log:
    names = subprocess.check_output(['git','ls-tree','-r','--name-only',BASE,'src'],text=True).splitlines()
    sources = [x for x in names if x.endswith('.c') and Path(x).name not in ('main.c', 'ray.c', 'spacetime_minkowski.c')]
    # Pin all sources/headers to the baseline for reproducibility after changes.
    srcdir = a.work_dir/'baseline'; srcdir.mkdir(exist_ok=True)
    for name in names:
        if name.endswith(('.c','.h')):
            (srcdir/Path(name).name).write_text(source(name))
    sources = [str(srcdir/Path(x).name) for x in sources]
    build(['cc', '-std=c11', '-march=native', '-pipe', '-Wall', '-Wextra', '-Wpedantic', '-O2', '-DNDEBUG', '-fopenmp', '-DENABLE_PNG', '-DFRAME_PSF_EVENT_SINK=1', '-DSPACETIME_SCHWARZSCHILD', '-I'+str(srcdir), *sources, str(a.work_dir/'ray.c'), str(a.work_dir/'main.c'), '-lm', '-lpng', '-o', str(a.work_dir/'renderer')])
(HERE/(a.tag+'-build-command.txt')).write_text('\n'.join(commands)+'\n')
metadata = {'base': BASE, 'cpu': subprocess.check_output(['lscpu'],text=True), 'compiler':subprocess.check_output(['cc','--version'],text=True), 'track_sha256':hashlib.sha256((HERE/'camera.csv').read_bytes()).hexdigest(), 'catalog_sha256':hashlib.sha256((ROOT/'assets/sky_grid_5deg.csv').read_bytes()).hexdigest()}
metadata.update(production=a.production, ray_sha256=hashlib.sha256(ray.encode()).hexdigest(),
                openmp_environment={k:v for k,v in os.environ.items() if k.startswith(('OMP_', 'GOMP_'))})
(HERE/(a.tag+'-environment.json')).write_text(json.dumps(metadata,indent=2))
results=[]
for number, schedule in enumerate(a.schedules):
    name = f'{a.tag}-{number}-{schedule.replace(",", "-")}'
    out=a.work_dir/name; out.mkdir(exist_ok=True)
    args=[str(a.work_dir/'renderer'), '--catalog','assets/sky_grid_5deg.csv', '--observer-track',str(HERE/'camera.csv'), '--movie-track-samples','--frames-dir',str(out), '--frames-prefix','frame', '--width','64','--height','36','--fov-deg','60','--exposure','0.2','--coarse-cell-pixels','16','--refine-max-level','3','--lens-map-output',str(out/'movie.grlens'),'--verbose']
    overrides=dict(OMP_NUM_THREADS='16', OMP_DYNAMIC='FALSE', OMP_SCHEDULE=schedule)
    command='env '+' '.join(shlex.quote(k+'='+v) for k,v in overrides.items())+' '+shlex.join(args)
    (HERE/(name+'-command.txt')).write_text(command+'\n')
    before=resource.getrusage(resource.RUSAGE_CHILDREN); start=time.monotonic()
    with (out/'raw.log').open('w') as log:
        proc=subprocess.run(args,env=dict(os.environ,**overrides),stdout=log,stderr=subprocess.STDOUT)
    elapsed=time.monotonic()-start; after=resource.getrusage(resource.RUSAGE_CHILDREN)
    raw=(out/'raw.log').read_bytes()
    with gzip.open(HERE/(name+'.log.gz'),'wb') as log: log.write(raw)
    row=dict(name=name,schedule=schedule,exit=proc.returncode,elapsed=elapsed,user=after.ru_utime-before.ru_utime,system=after.ru_stime-before.ru_stime)
    if proc.returncode==0:
        row['lens_map_sha256']=hashlib.sha256((out/'movie.grlens').read_bytes()).hexdigest()
        row['png_count']=len(list(out.glob('frame_*.png')))
    results.append(row)
    (HERE/(a.tag+'-results.json')).write_text(json.dumps(results,indent=2))
    print(json.dumps(row),flush=True)
    if proc.returncode: raise SystemExit(proc.returncode)