// gpu-bench: times every model mcap runs on the GPU (each OCR model, the label model's image encoder,
// segmentation) over a fixed corpus, so machines can be compared. The corpus (gpu-bench-corpus beside
// this program, from tools/gpu_bench/corpus) is checked file by file against corpus.txt's SHA-256s
// before anything runs, and each model's files and models.config entry are fingerprinted, so a
// comparison of two results is known to be like for like. mcap's GPU Benchmark window runs this and
// lays the results side by side.
//
//   gpu-bench [-only ocr,labels,segment] [-models a,b] [-repeats N] [-warmup N] [-cpu | -cpu-only]
//             [-corpus FOLDER] [-out FOLDER] [-note TEXT] [-progress]
//
//   -only      the kinds of model to run (default all)
//   -models    the models to run, by their models.config names (default every one)
//   -repeats   timed passes over the corpus; the median is reported (default 3)
//   -warmup    untimed passes first, so every image size has been seen (default 1)
//   -cpu       the CPU too, after the GPU (slow: minutes for segmentation); also times the label
//              model's text encoder, which always runs on the CPU, as a like: search uses it
//   -cpu-only  the CPU only
//   -corpus    the corpus folder (default gpu-bench-corpus beside this program)
//   -out       where the results file goes (default gpu-bench-results beside this program)
//   -note      a few words kept with the results, such as "plugged in, max performance"
//   -progress  also print @progress lines, for the GPU Benchmark window
//   -ocr-workers N  OCR's worker threads on the GPU, in place of Ocr_Options' (for investigating memory)
//   -ort-log LEVEL  onnxruntime's console log: verbose, info (each GPU arena extension), warning (default)
//
// Build: jai first.jai - release (it is built with mcap on the gpu-bench branch).

USAGE :: "Usage: gpu-bench [-only ocr,labels,segment] [-models a,b] [-repeats N] [-warmup N] [-cpu | -cpu-only] [-corpus FOLDER] [-out FOLDER] [-note TEXT] [-progress] [-ocr-workers N] [-ort-log LEVEL]\n";

options: struct {
  only: [..] string;
  models: [..] string;
  repeats := 3;
  warmup := 1;
  gpu := true;
  cpu := false;
  progress := false;
  ocr_workers: s32;   // 0: Ocr_Options' default
  corpus_folder, out_folder, note: string;
};

main :: () {
  runtime_support_headless = true;   // a command-line tool: a failure is reported on the console and exits with 1, no message box
  args := get_command_line_arguments();
  exe_folder := path_strip_filename(get_path_of_running_executable());
  options.corpus_folder = sprint("%gpu-bench-corpus", exe_folder);
  options.out_folder    = sprint("%gpu-bench-results", exe_folder);
  i := 1;
  while i < args.count {
    arg := args[i];
    value := ifx i + 1 < args.count then args[i + 1] else "";
    if arg == {
      case "-only";     for split(value, ",")  array_add(*options.only, trim(it)); i += 2;
      case "-models";   for split(value, ",")  array_add(*options.models, trim(it)); i += 2;
      case "-repeats";  options.repeats = max(1, string_to_int(value)); i += 2;
      case "-warmup";   options.warmup = max(0, string_to_int(value)); i += 2;
      case "-cpu";      options.cpu = true; i += 1;
      case "-cpu-only"; options.cpu = true; options.gpu = false; i += 1;
      case "-corpus";   options.corpus_folder = value; i += 2;
      case "-out";      options.out_folder = value; i += 2;
      case "-note";     options.note = value; i += 2;
      case "-progress"; options.progress = true; i += 1;
      case "-ocr-workers"; options.ocr_workers = cast(s32) max(1, string_to_int(value)); i += 2;
      case "-ort-log";
        if value == {
          case "verbose"; ort.log_level = .ORT_LOGGING_LEVEL_VERBOSE;
          case "info";    ort.log_level = .ORT_LOGGING_LEVEL_INFO;
          case "warning"; ort.log_level = .ORT_LOGGING_LEVEL_WARNING;
          case;           print("-ort-log takes verbose, info or warning.\n"); exit(2);
        }
        i += 2;
      case "-help";     print(USAGE); exit(0);
      case;             print("Unknown option %.\n%", arg, USAGE); exit(2);
    }
  }
  make_directory_if_it_does_not_exist(options.out_folder, recursive = true);
  started_at := current_time_consensus();

  // The machine.
  info := system_info();
  ort_ok, ort_error := ort_load();
  print("gpu-bench, mcap % (%, %)\n", VERSION, GIT_HASH, GIT_BRANCH);
  print("machine:     %, %\n", info.machine, info.os);
  print("cpu:         % (% threads), % GB of memory\n", info.cpu, info.logical_cores, FF(info.ram_gb, 1));
  if info.gpu  print("gpu:         %, % GB, driver % (CUDA %), power limit % W, max clocks % / % MHz\n", info.gpu, FF(info.vram_gb, 1), info.driver, info.cuda_driver, FF(info.power_limit_w, 0), info.max_sm_clock_mhz, info.max_mem_clock_mhz);
  else         print("gpu:         none (%)\n", info.gpu_error);
  if info.power_source  print("power:       %\n", info.power_source);
  if !ort_ok { print("%\n", ort_error); exit(1); }
  print("onnxruntime: % (%)\n", ort.version, ort.library_path);

  // The corpus: every file as corpus.txt pins it, or nothing runs.
  progress("corpus", "-", "checking", 0, 0);
  corpus, corpus_error := corpus_load(options.corpus_folder);
  if corpus_error { print("\nThe corpus cannot be used: %\n", corpus_error); exit(1); }
  print("corpus:      % images, % megapixels, fingerprint % (decoded in % s)\n\n", corpus.images.count, FF(corpus.megapixels, 1), short(corpus.fingerprint), FF(corpus.decode_s, 2));

  // The models, and what each is made of.
  tasks := list_tasks(exe_folder);
  if !tasks.count { print("No models to run: none matched -only and -models, or none are beside this program (run jai first.jai).\n"); exit(1); }
  hash_cache_load();
  for * tasks  fingerprint_task(it);
  hash_cache_save();

  runs: [..] Run;
  for * task: tasks {
    providers: [..] string;
    providers.allocator = temp;
    if task.kind == .LABELS_TEXT { if options.cpu  array_add(*providers, "CPU"); }
    else {
      if options.gpu  array_add(*providers, "CUDA");
      if options.cpu  array_add(*providers, "CPU");
    }
    for provider: providers {
      run := run_task(task, provider == "CUDA", *corpus);
      array_add(*runs, run);
      if !run.ok { print("== % on %: %\n\n", task.name, provider, run.error); continue; }
      print_run(task, run);
    }
  }

  cudart, cudnn := cuda_library_versions();
  path := write_results(info, started_at, *corpus, tasks, runs, cudart, cudnn);
  print("Results: %\n", path);
  if options.progress  print("@results %\n", path);
}

Task_Kind :: enum { OCR; LABELS; LABELS_TEXT; SEGMENT; }

Task :: struct {
  kind: Task_Kind;
  model: string;           // as models.config names it
  name: string;            // kind/model: ocr/ppocr-v5-mobile
  models_folder: string;
  files: [..] string;      // the model's files, in models_folder
  config: string;          // its models.config entry, as compared
  fingerprint: string;     // of files and config
  pinned: string;          // "yes" when every file is one a fetch script pins, "no" when one differs, "-" when none are pinned
}

Run :: struct {
  task: *Task;
  ok: bool;
  error: string;
  provider: string;        // what it really ran on: a GPU run can fall back to the CPU
  items: s64;              // per pass
  unit: string;            // image or phrase
  megapixels: float64;     // per pass
  load_s, first_s: float64;
  pass_s: [..] float64;
  output: string;          // what the models found, in words without spaces: lines:123,chars:4567
  digest: string;          // of what they found, to see whether two machines' models agree
  stages: string;          // where the time went, per item
  batch: s64;
  gpu: Gpu_Stats;
  budget: string;          // OCR on the GPU: what its memory budget measured (modules/OCR/budget.jai), or ""
  trace: [..] string;      // the GPU's memory above the baseline, MB, after each step: load:812 first:1020 warm1:... pass1:.../peak...
  explanation: string;     // why a GPU run is on the CPU
}

Gpu_Stats :: struct {
  sampled: bool;
  baseline_mb, peak_mb: float64;   // device memory before the engine, and the most during the timed passes
  pass_peak_mb: float64;           // the most during the current timed pass
  util_sum, power_sum, clock_sum: float64;
  temperature_max: float64;
  samples: s64;
}

// The corpus, decoded.
Corpus :: struct {
  fingerprint: string;
  images: [..] Corpus_Image;
  megapixels: float64;
  decode_s: float64;
}

Corpus_Image :: struct { name: string; pixels: *u8; width, height: s32; }

// Every file corpus.txt names, checked against its hash; the error says what is wrong.
corpus_load :: (folder: string) -> Corpus, error: string {
  corpus: Corpus;
  manifest_path := tprint("%/corpus.txt", folder);
  manifest, ok := read_entire_file(manifest_path, log_errors = false);
  if !ok  return corpus, tprint("% cannot be read. The build copies it, with the images, from tools/gpu_bench.", manifest_path);
  fingerprint: Sha256;
  bad: [..] string;
  bad.allocator = temp;
  start := current_time_monotonic();
  for raw: split(manifest, "\n") {
    line := trim(raw);
    if !line.count || line[0] == #char "#"  continue;
    space := find_index_from_left(line, #char " ");
    if space < 0  continue;
    expected := slice(line, 0, space);
    name := trim(advance(line, space));
    path := tprint("%/%", folder, name);
    data, read_ok := read_entire_file(path, log_errors = false);
    if !read_ok { array_add(*bad, tprint("% is missing", name)); continue; }
    hash := sha256_of_string(data,, temp);
    if hash != expected { array_add(*bad, tprint("% is not the file corpus.txt pins (SHA-256 %)", name, hash)); free(data); continue; }
    sha256_update(*fingerprint, tprint("% %\n", name, hash));
    progress("corpus", "-", "decoding", corpus.images.count + 1, 0);
    image, decoded, decode_error := image_decode(cast([] u8) data);
    free(data);
    if !decoded { array_add(*bad, tprint("% cannot be decoded (%)", name, decode_error)); continue; }
    array_add(*corpus.images, .{copy_string(name), image.pixels, image.width, image.height});
    corpus.megapixels += cast(float64) image.width * image.height / 1_000_000;
    reset_temporary_storage();
  }
  corpus.decode_s = seconds_since(start);
  if bad.count  return corpus, join(..bad, "; ");
  if !corpus.images.count  return corpus, tprint("% names no images.", manifest_path);
  corpus.fingerprint = sha256_finish(*fingerprint);
  return corpus, "";
}

// The models beside this program, as -only and -models choose.
list_tasks :: (exe_folder: string) -> [..] Task {
  tasks: [..] Task;
  add :: (tasks: *[..] Task, kind: Task_Kind, model: string, folder: string, files: .. string) {
    prefix := ifx kind == .OCR then "ocr" else ifx kind == .LABELS then "labels" else ifx kind == .LABELS_TEXT then "labels-text" else "segment";
    if options.only.count && !array_find(options.only, ifx kind == .LABELS_TEXT then "labels" else prefix)  return;
    if options.models.count && !array_find(options.models, model)  return;
    task := array_add(tasks);
    task.* = .{kind = kind, model = model, name = sprint("%/%", prefix, model), models_folder = folder};
    for files  if it  array_add(*task.files, it);
    task.config = config_entry(tprint("%/models.config", folder), model);
  }
  ocr_folder := sprint("%ocr-models", exe_folder);
  ocr_models, ocr_error := ocr_list_models(ocr_folder);
  if ocr_error  print("OCR: %\n", ocr_error);
  for ocr_models  add(*tasks, .OCR, it.name, ocr_folder, it.detector, it.recognizer, it.dictionary, it.textline_ori);
  labels_folder := sprint("%label-models", exe_folder);
  labels_models, labels_error := labels_list_models(labels_folder);
  if labels_error  print("Labels: %\n", labels_error);
  for labels_models  add(*tasks, .LABELS, it.name, labels_folder, it.vision, it.text, it.tokenizer);
  for labels_models  add(*tasks, .LABELS_TEXT, it.name, labels_folder, it.vision, it.text, it.tokenizer);
  segment_folder := sprint("%segment-models", exe_folder);
  segment_models, segment_error := segment_list_models(segment_folder);
  if segment_error  print("Segmentation: %\n", segment_error);
  for segment_models  add(*tasks, .SEGMENT, it.name, segment_folder, it.model, it.classes);
  return tasks;
}

// A model's [model NAME] entry, comments and spacing taken out, so only a change that matters
// changes the fingerprint.
config_entry :: (path: string, model: string) -> string {
  text, ok := read_entire_file(path, log_errors = false,, temp);
  if !ok  return "";
  builder: String_Builder;
  inside := false;
  for raw: split(text, "\n",, temp) {
    line := raw;
    hash := find_index_from_left(line, #char "#");
    if hash >= 0  line = slice(line, 0, hash);
    line = trim(line);
    if !line.count  continue;
    if begins_with(line, "[model ") { inside = line == tprint("[model %]", model); continue; }
    if !inside  continue;
    colon := find_index_from_left(line, #char ":");
    if colon >= 0  line = tprint("%: %", trim(slice(line, 0, colon)), trim(advance(line, colon + 1)));
    print_to_builder(*builder, "%\n", line);
  }
  return builder_to_string(*builder);
}

// The files the fetch scripts pin (modules/OCR/fetch_models.jai, modules/Labels/fetch_models.jai).
// Segmentation's model is made on each machine by tools/export_mask2former.py, so it has no pin:
// copy the .onnx from one machine to the others to be sure they run the same one.
Pin :: struct { name: string; sha256: string; }
PINNED :: Pin.[
  .{"ch_PP-OCRv5_det_mobile.onnx", "4d97c44a20d30a81aad087d6a396b08f786c4635742afc391f6621f5c6ae78ae"},
  .{"ch_PP-OCRv5_rec_mobile.onnx", "5825fc7ebf84ae7a412be049820b4d86d77620f204a041697b0494669b1742c5"},
  .{"ppocrv5_dict.txt",            "d1979e9f794c464c0d2e0b70a7fe14dd978e9dc644c0e71f14158cdf8342af1b"},
  .{"ch_PP-LCNet_x0_25_textline_ori_cls_mobile.onnx", "54379ae5174d026780215fc748a7f31910dee36818e63d49e17dc598ecc82df7"},
  .{"ch_PP-OCRv5_det_server.onnx", "0f8846b1d4bba223a2a2f9d9b44022fbc22cc019051a602b41a7fda9667e4cad"},
  .{"ch_PP-OCRv5_rec_server.onnx", "e09385400eaaaef34ceff54aeb7c4f0f1fe014c27fa8b9905d4709b65746562a"},
  .{"siglip2-base-patch16-384-vision.onnx",    "49ae4958b1098ca995e929d646f7be05a69c65e6344beae07d58c6598ffc5210"},
  .{"siglip2-base-patch16-384-text-int8.onnx", "30a94b43b9ddc39c04ab67f80cba4b808262682b2fa1a2a8b924bf1e7bd84c61"},
  .{"siglip2-tokenizer.json",                  "cb9140fae3ac5122c972d37adf83e1248471a38147ad76f8215c8872c6fd8322"},
];

fingerprint_task :: (task: *Task) {
  h: Sha256;
  pinned, differs := 0, 0;
  for task.files {
    progress("hashing", task.name, it, it_index + 1, task.files.count);
    hash := hash_file_cached(tprint("%/%", task.models_folder, it));
    if !hash  hash = "missing";
    sha256_update(*h, tprint("% %\n", it, hash));
    for pin: PINNED  if pin.name == it { if pin.sha256 == hash  pinned += 1; else differs += 1; }
  }
  sha256_update(*h, task.config);
  task.fingerprint = sha256_finish(*h);
  task.pinned = ifx differs then "no" else ifx pinned then "yes" else "-";
}

run_task :: (task: *Task, prefer_gpu: bool, corpus: *Corpus) -> Run {
  run: Run;
  run.task = task;
  run.provider = ifx prefer_gpu then "CUDA" else "CPU";
  run.unit = ifx task.kind == .LABELS_TEXT then "phrase" else "image";
  images := corpus.images;
  run.items = ifx task.kind == .LABELS_TEXT then PHRASES.count else images.count;
  if task.kind != .LABELS_TEXT  run.megapixels = corpus.megapixels;

  sampler_start(*run.gpu);   // the baseline: the device's memory before the engine
  sampler_pause();

  // The engine.
  progress(task.name, run.provider, "loading", 0, 0);
  load_start := current_time_monotonic();
  engine: struct { ocr: *Ocr_Engine; labels: *Labels_Engine; segment: *Segment_Engine; };
  ok: bool;
  error: string;
  phrase_vectors: [] float;
  if task.kind == {
    case .OCR;
      ocr_options: Ocr_Options;
      if options.ocr_workers  ocr_options.gpu_workers = options.ocr_workers;
      engine.ocr, ok, error = ocr_create(task.models_folder, task.model, prefer_gpu, ocr_options);
      if ok {
        ocr_gpu, ocr_why := ocr_gpu_status(engine.ocr);
        if prefer_gpu && !ocr_gpu  run.explanation = copy_string(ocr_why);
        run.provider = sprint("%", ocr_provider(engine.ocr));
        run.batch = engine.ocr.rec_batch;
        b := engine.ocr.budget;
        if b.enabled  run.budget = sprint("detector %+% MB/MP, recognizer %+% MB per 1000 line-px", FF(b.det_fixed / MEGABYTE, 0), FF(b.det_per_pixel * 1_000_000 / MEGABYTE, 0),
                                         FF(b.rec_fixed / MEGABYTE, 0), FF(b.rec_per_column * 1000 / MEGABYTE, 1));
        else if engine.ocr.provider == .CUDA  run.budget = "off (the arena could not be measured)";
      }
    case .LABELS; #through;
    case .LABELS_TEXT;
      engine.labels, ok, error = labels_create(task.models_folder, task.model, prefer_gpu && task.kind == .LABELS);
      if ok {
        labels_gpu, labels_why := labels_gpu_status(engine.labels);
        if prefer_gpu && !labels_gpu  run.explanation = copy_string(labels_why);
        run.provider = ifx task.kind == .LABELS_TEXT then "CPU" else sprint("%", engine.labels.provider);
        run.batch = ifx task.kind == .LABELS_TEXT then 1 else engine.labels.model.image_batch;
        phrases := PHRASES;
        phrase_vectors = labels_embed_texts(engine.labels, phrases);
      }
    case .SEGMENT;
      engine.segment, ok, error = segment_create(task.models_folder, task.model, prefer_gpu);
      if ok {
        if prefer_gpu && engine.segment.provider != .CUDA  run.explanation = copy_string(engine.segment.gpu_explanation);
        run.provider = sprint("%", engine.segment.provider);
        run.batch = ifx engine.segment.provider == .CUDA then engine.segment.model.batch_gpu else engine.segment.model.batch_cpu;
      }
  }
  run.load_s = seconds_since(load_start);
  trace(*run, "load");
  if !ok { sampler_stop(); run.error = copy_string(error); return run; }
  defer {
    ocr_destroy(engine.ocr);
    labels_destroy(engine.labels);
    segment_destroy(engine.segment);
    array_free(phrase_vectors);
  }
  if prefer_gpu && run.provider != "CUDA" {
    // Said, and not timed: the CPU run, when asked for, is the same thing.
    sampler_stop();
    run.error = tprint("not on the GPU: %", run.explanation,, context.allocator);
    return run;
  }

  // One pass over the corpus, or over its first image alone; what it found, and its digest.
  pass :: (first_only: bool) -> output: string, digest: string #expand {
    h: Sha256;
    builder: String_Builder;
    builder.allocator = temp;
    count := ifx first_only then 1 else `run.items;
    if `task.kind == {
      case .OCR;
        lines, characters: s64;
        done := 0;
        while done < count {
          n := min(OCR_BATCH, count - done);
          batch := NewArray(n, Ocr_Image,, temp);
          for 0..n - 1 { image := `images[done + it]; batch[it] = .{image.pixels, image.width, image.height, image.width * 4}; }
          results := ocr_process(`engine.ocr, batch,, temp);
          for results {
            if it.failed { sha256_update(*h, "failed\n"); continue; }
            lines += it.lines.count;
            for line: it.lines { characters += line.text.count; sha256_update(*h, line.text); sha256_update(*h, "\n"); }
          }
          done += n;
        }
        print_to_builder(*builder, "lines:%,chars:%", lines, characters);
      case .LABELS;
        input := NewArray(count, Labels_Image,, temp);
        for 0..count - 1 { image := `images[it]; input[it] = .{image.pixels, image.width, image.height, image.width * 4}; }
        vectors, error := labels_embed_images(`engine.labels, input,, temp);
        if error  sha256_update(*h, "failed\n");
        // Each image's best phrase: what the label list would make of it.
        dims := `engine.labels.dimensions;
        best_count: [PHRASES.count] s64;
        for image: 0..count - 1 {
          best, best_score := -1, -1.0;
          for phrase: 0..PHRASES.count - 1 {
            score := labels_score(`engine.labels, vectors.data + image * dims, `phrase_vectors.data + phrase * dims);
            if score > best_score { best, best_score = phrase, score; }
          }
          sha256_update(*h, tprint("%\n", best));
          best_count[best] += 1;
        }
        distinct := 0;
        for best_count  if it  distinct += 1;
        print_to_builder(*builder, "best_phrases:%", distinct);
      case .LABELS_TEXT;
        for 0..count - 1 {
          one: [1] string;
          one[0] = PHRASES[it];
          vectors, error := labels_embed_texts(`engine.labels, one,, temp);   // one at a time, as a search does
          sha256_update(*h, ifx error then "failed\n" else "ok\n");
        }
        print_to_builder(*builder, "phrases:%", count);
      case .SEGMENT;
        input := NewArray(count, Segment_Image,, temp);
        for 0..count - 1 { image := `images[it]; input[it] = .{image.pixels, image.width, image.height, image.width * 4}; }
        results := segment_images(`engine.segment, input,, temp);
        segments: s64;
        for results {
          if it.failed { sha256_update(*h, "failed\n"); continue; }
          segments += it.segments.count;
          for segment: it.segments  sha256_update(*h, tprint("% ", segment.class));
          sha256_update(*h, "\n");
        }
        print_to_builder(*builder, "segments:%", segments);
    }
    return builder_to_string(*builder,, temp), sha256_finish(*h,, temp);
  }

  // The first image alone: the first run pays for choosing kernels and setting up memory.
  progress(task.name, run.provider, "first", 0, 0);
  first_start := current_time_monotonic();
  pass(true);
  run.first_s = seconds_since(first_start);
  trace(*run, "first");
  reset_temporary_storage();
  for 1..options.warmup {
    progress(task.name, run.provider, "warm-up", it, options.warmup);
    pass(false);
    trace(*run, tprint("warm%", it));
    reset_temporary_storage();
  }
  if task.kind == .OCR      ocr_reset_stats(engine.ocr);
  if task.kind == .SEGMENT  engine.segment.stats = .{};

  sampler_resume();
  for 1..options.repeats {
    progress(task.name, run.provider, "timing", it, options.repeats);
    start := current_time_monotonic();
    sampler_pass_start();
    output, digest := pass(false);
    array_add(*run.pass_s, seconds_since(start));
    trace(*run, tprint("pass%", it), sampler_pass_peak());
    if it == 1 { run.output = copy_string(output); run.digest = copy_string(short(digest)); }
    reset_temporary_storage();
  }
  sampler_stop();

  per := cast(float64) max(run.items * options.repeats, 1) * 1000;
  if task.kind == .OCR {
    s := ocr_stats(engine.ocr);
    run.stages = sprint("detect:%,orient:%,recognize:%,budget_waits:%", FF(s.detect_us / per, 2), FF(s.orientation_us / per, 2), FF(s.recognize_us / per, 2), s.budget_waits);
    if engine.ocr.budget.enabled  run.budget = sprint("%; limit % MB at the end", run.budget, FF(engine.ocr.budget.limit / MEGABYTE, 0));
  } else if task.kind == .SEGMENT {
    s := engine.segment.stats;
    run.stages = sprint("prepare:%,model:%,post:%", FF(s.prepare_us / per, 2), FF(s.model_us / per, 2), FF(s.post_us / per, 2));
  }
  run.ok = true;
  return run;
}

OCR_BATCH :: 8;   // images per ocr_process call, as the app's scan gives them

// Phrases for the label model: its image vectors are scored against them (the output's digest), and
// the text encoder is timed on them, one at a time, as a like: search runs it.
PHRASES :: string.[
  "a screenshot of a code editor showing source code", "a screenshot of a terminal window",
  "a spreadsheet with rows and columns of numbers", "a document with paragraphs of text",
  "a chat conversation on a phone", "a dialog box asking to save changes",
  "a photo of a living room", "a photo of people playing tennis",
  "a photo of food on a plate", "a photo of a street with cars",
  "a photo of an animal", "a photo of a person", "an electronic circuit schematic diagram",
  "a printed circuit board", "a graph with several colored lines", "a web page in a browser",
];

Median :: struct { median, best, worst: float64; }

median_of :: (values: [] float64) -> Median {
  sorted := array_copy(values,, temp);
  quick_sort(sorted, (a: float64, b: float64) -> s64 { return ifx a < b then -1 else ifx a > b then 1 else 0; });
  m: Median;
  if !sorted.count  return m;
  m.median = ifx sorted.count % 2 then sorted[sorted.count / 2] else (sorted[sorted.count / 2 - 1] + sorted[sorted.count / 2]) / 2;
  m.best, m.worst = sorted[0], sorted[sorted.count - 1];
  return m;
}

print_run :: (task: *Task, run: Run) {
  m := median_of(run.pass_s);
  print("== % on % (batch %)\n", task.name, run.provider, run.batch);
  print("   load % s, first % % s\n", FF(run.load_s, 2), run.unit, FF(run.first_s, 3));
  print("   % ms per % (median of %; best %, worst %), % %s/s", FF(1000 * m.median / run.items, 2), run.unit, run.pass_s.count,
        FF(1000 * m.best / run.items, 2), FF(1000 * m.worst / run.items, 2), FF(run.items / m.median, 1), run.unit);
  if run.megapixels  print(", % megapixels/s", FF(run.megapixels / m.median, 1));
  print("\n");
  if run.stages  print("   per %, summed over threads (ms): %\n", run.unit, run.stages);
  if run.budget  print("   GPU memory budget: %\n", run.budget);
  if run.gpu.samples  print("   GPU: % MB more memory at the peak, % percent busy, % W, % MHz, up to % C\n",
                            FF(run.gpu.peak_mb - run.gpu.baseline_mb, 0), FF(run.gpu.util_sum / run.gpu.samples, 0),
                            FF(run.gpu.power_sum / run.gpu.samples, 0), FF(run.gpu.clock_sum / run.gpu.samples, 0), FF(run.gpu.temperature_max, 0));
  if run.trace.count  print("   GPU memory over the baseline (MB), after each step [and its peak]: %\n", join(..run.trace, " ",, temp));
  print("   found % (digest %)\n\n", run.output, run.digest);
}

write_results :: (info: System_Info, started_at: Apollo_Time, corpus: *Corpus, tasks: [] Task, runs: [] Run, cudart: string, cudnn: string) -> path: string {
  b: String_Builder;
  print_to_builder(*b, "%\n", BENCH_RESULTS_FORMAT);
  date := to_calendar(started_at, .LOCAL);
  print_to_builder(*b, "date: %-%-% %:%:%\n", date.year, pad2(date.month_starting_at_0 + 1), pad2(date.day_of_month_starting_at_0 + 1), pad2(date.hour), pad2(date.minute), pad2(date.second));
  print_to_builder(*b, "machine: %\n", one_line(info.machine));
  if options.note  print_to_builder(*b, "note: %\n", one_line(options.note));
  print_to_builder(*b, "os: %\n", info.os);
  print_to_builder(*b, "cpu: % (% threads)\n", one_line(info.cpu), info.logical_cores);
  print_to_builder(*b, "ram_gb: %\n", FF(info.ram_gb, 1));
  print_to_builder(*b, "gpu: %\n", ifx info.gpu then one_line(info.gpu) else tprint("none (%)", info.gpu_error));
  if info.gpu {
    print_to_builder(*b, "vram_gb: %\n", FF(info.vram_gb, 1));
    print_to_builder(*b, "driver: % (CUDA %)\n", info.driver, info.cuda_driver);
    print_to_builder(*b, "power_limit_w: %\n", FF(info.power_limit_w, 0));
    print_to_builder(*b, "max_clocks_mhz: % / %\n", info.max_sm_clock_mhz, info.max_mem_clock_mhz);
  }
  if info.power_source  print_to_builder(*b, "power: %\n", info.power_source);
  print_to_builder(*b, "onnxruntime: %\n", ort.version);
  if cudart  print_to_builder(*b, "cudart: %\n", cudart);
  if cudnn   print_to_builder(*b, "cudnn: %\n", cudnn);
  print_to_builder(*b, "mcap: % (%, %)\n", VERSION, GIT_HASH, GIT_BRANCH);
  print_to_builder(*b, "passes: % timed, % warm-up\n", options.repeats, options.warmup);
  print_to_builder(*b, "corpus: %\n", corpus.fingerprint);
  print_to_builder(*b, "corpus_size: % images, % megapixels\n", corpus.images.count, FF(corpus.megapixels, 1));
  for tasks {
    print_to_builder(*b, "model: % %\n", it.name, it.fingerprint);
    print_to_builder(*b, "pinned %: %\n", it.name, it.pinned);
  }
  for run: runs {
    task := run.task;
    if !run.ok { print_to_builder(*b, "error % %: %\n", task.name, run.provider, one_line(run.error)); continue; }
    m := median_of(run.pass_s);
    print_to_builder(*b, "result: % % unit=% items=% batch=% ms=% best_ms=% worst_ms=% per_s=% load_s=% first_s=%",
                     task.name, run.provider, run.unit, run.items, run.batch, FF(1000 * m.median / run.items, 3), FF(1000 * m.best / run.items, 3),
                     FF(1000 * m.worst / run.items, 3), FF(run.items / m.median, 2), FF(run.load_s, 2), FF(run.first_s, 3));
    if run.megapixels  print_to_builder(*b, " mp_per_s=%", FF(run.megapixels / m.median, 2));
    if run.gpu.samples {
      g := run.gpu;
      print_to_builder(*b, " vram_mb=% gpu_util=% power_w=% sm_mhz=% temp_c=%", FF(g.peak_mb - g.baseline_mb, 0), FF(g.util_sum / g.samples, 0),
                       FF(g.power_sum / g.samples, 0), FF(g.clock_sum / g.samples, 0), FF(g.temperature_max, 0));
    }
    if run.stages  print_to_builder(*b, " stages=%", run.stages);
    if run.trace.count  print_to_builder(*b, " trace=%", join(..run.trace, ",",, temp));
    print_to_builder(*b, " output=% digest=%\n", run.output, run.digest);
  }
  name := tprint("%_%-%-%_%-%-%.txt", file_safe(info.machine), date.year, pad2(date.month_starting_at_0 + 1), pad2(date.day_of_month_starting_at_0 + 1),
                 pad2(date.hour), pad2(date.minute), pad2(date.second));
  path := tprint("%/%", options.out_folder, name);
  if !write_entire_file(path, *b)  print("Could not write %.\n", path);
  return path;
}

// For the GPU Benchmark window: where the run is. total 0: no count.
progress :: (task: string, provider: string, phase: string, done: s64, total: s64) {
  if !options.progress  return;
  print("@progress % % % % %\n", task, provider, phase, done, total);
}

// The GPU, sampled every SAMPLE_MS on a thread of its own while the timed passes run.
SAMPLE_MS :: 100;

sampler: struct {
  thread: Thread;
  stats: *Gpu_Stats;
  running: bool;   // the thread is up
  paused: bool;    // it is up, and not sampling
  mutex: Mutex;
};

sampler_start :: (stats: *Gpu_Stats) {
  s, ok := gpu_sample();
  if !ok  return;
  stats.sampled = true;
  stats.baseline_mb = s.memory_used_mb;
  sampler.stats = stats;
  init(*sampler.mutex);
  sampler.running = true;
  sampler.paused = false;
  thread_init(*sampler.thread, (thread: *Thread) -> s64 {
    while sampler.running {
      sleep_milliseconds(SAMPLE_MS);
      s, ok := gpu_sample();
      lock(*sampler.mutex);
      if ok && !sampler.paused && sampler.running {
        g := sampler.stats;
        g.peak_mb = max(g.peak_mb, s.memory_used_mb);
        g.pass_peak_mb = max(g.pass_peak_mb, s.memory_used_mb);
        g.util_sum += s.util_percent;
        g.power_sum += s.power_w;
        g.clock_sum += s.sm_clock_mhz;
        g.temperature_max = max(g.temperature_max, s.temperature_c);
        g.samples += 1;
      }
      unlock(*sampler.mutex);
    }
    return 0;
  });
  thread_start(*sampler.thread);
}

// The GPU's memory now, over the run's baseline, as a step of its trace; with the step's peak when known.
trace :: (run: *Run, step: string, peak: float64 = -1) {
  if !run.gpu.sampled  return;
  s, ok := gpu_sample();
  if !ok  return;
  now := s.memory_used_mb - run.gpu.baseline_mb;
  if peak >= 0  array_add(*run.trace, sprint("%:%[%]", step, FF(now, 0), FF(peak - run.gpu.baseline_mb, 0)));
  else          array_add(*run.trace, sprint("%:%", step, FF(now, 0)));
}

// The peak of one timed pass: reset as it starts, read as it ends.
sampler_pass_start :: () { if !sampler.running  return; lock(*sampler.mutex); sampler.stats.pass_peak_mb = 0; unlock(*sampler.mutex); }
sampler_pass_peak :: () -> float64 {
  if !sampler.running  return -1;
  lock(*sampler.mutex);
  peak := sampler.stats.pass_peak_mb;
  unlock(*sampler.mutex);
  return peak;
}

sampler_pause  :: () { if !sampler.running  return; lock(*sampler.mutex); sampler.paused = true;  unlock(*sampler.mutex); }
sampler_resume :: () { if !sampler.running  return; lock(*sampler.mutex); sampler.paused = false; unlock(*sampler.mutex); }

sampler_stop :: () {
  if !sampler.running  return;
  lock(*sampler.mutex);
  sampler.running = false;
  unlock(*sampler.mutex);
  thread_is_done(*sampler.thread, -1);
  thread_deinit(*sampler.thread);
  destroy(*sampler.mutex);
}

// Model files are hundreds of MB, so their hashes are kept in the results folder, by path, size and
// time: hashed again only when one changes.
hash_cache: [..] struct { key: string; hash: string; };

hash_cache_load :: () {
  text, ok := read_entire_file(tprint("%/.hash-cache", options.out_folder), log_errors = false);
  if !ok  return;
  for split(text, "\n") {
    tab := find_index_from_left(it, #char "\t");
    if tab > 0  array_add(*hash_cache, .{trim(advance(it, tab + 1)), slice(it, 0, tab)});
  }
}

hash_cache_save :: () {
  b: String_Builder;
  for hash_cache  print_to_builder(*b, "%\t%\n", it.hash, it.key);
  write_entire_file(tprint("%/.hash-cache", options.out_folder), *b);
}

hash_file_cached :: (path: string) -> string {
  modtime, size, ok := file_modtime_and_size(path);
  if !ok  return "";
  key := sprint("% % % %", size, modtime.low, modtime.high, path);
  for hash_cache  if it.key == key  return it.hash;
  hash := sha256_of_file(path);
  if hash  array_add(*hash_cache, .{key, hash});
  return hash;
}

MEGABYTE :: 1024.0 * 1024;

seconds_since :: (start: Apollo_Time) -> float64 {
  return to_float64_seconds(current_time_monotonic() - start);
}

short :: (hash: string) -> string { return slice(hash, 0, min(hash.count, 12)); }

pad2 :: (n: s64) -> string { return tprint("%", formatInt(n, minimum_digits = 2)); }

one_line :: (text: string) -> string {
  return replace(replace(trim(text), "\r", " ",, temp), "\n", " ",, temp);
}

file_safe :: (text: string) -> string {
  result := copy_string(text,, temp);
  for 0..result.count - 1 {
    c := result[it];
    if !(is_alnum(c) || c == #char "-")  result[it] = #char "_";
  }
  return result;
}

#load "sha256.jai";
#load "system_info.jai";
#load "results.jai";

#import "Basic";
#import "String";
#import "File";
#import "File_Utilities";
#import "Sort";
#import "System";
#import "Thread";
#import "Image_Decode";
#import "onnxruntime";
#import "OCR";
#import "Labels";
#import "Segment";