// ocr-bench: runs OCR over a folder of .mcap captures with each model, on the GPU when it can and
// on the CPU, and reports the time per image and per stage. Each capture's lines are written to
// OUT/<model>-<provider>/<name>.txt, to read against the images.
//
//   ocr-bench MCAP_FOLDER [-models a,b] [-cpu | -gpu-only] [-batch N] [-limit N] [-out FOLDER]
//             [-workers N] [-rec-batch N] [-warmup N]
//
//   -models     the models to run (default: every model in ocr-models/models.config)
//   -cpu        the CPU only; by default the GPU too, when it is available
//   -gpu-only   the GPU only
//   -workers    on the GPU, threads preparing images and decoding results (default Ocr_Options')
//   -rec-batch  text lines per recognizer run, in place of models.config's
//   -warmup     untimed passes over every image before the timed one (default 0)
//   -batch    images per ocr_process call (default 8)
//   -limit    at most this many captures (default all)
//   -out      where the text goes (default ocr-out beside this program)
//
// The models are read from ocr-models beside this program, where the build copies them.

USAGE :: "Usage: ocr-bench MCAP_FOLDER [-models a,b] [-cpu | -gpu-only] [-batch N] [-limit N] [-out FOLDER] [-workers N] [-rec-batch N] [-warmup N]\n";

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();
  if args.count < 2 { print(USAGE); exit(2); }
  folder := args[1];
  model_names: [..] string;
  cpu_only := false;
  gpu_only := false;
  options: Ocr_Options;
  rec_batch: s32 = 0;   // 0: as models.config says
  warmup := 0;
  batch := 8;
  limit := S64_MAX_VALUE;
  exe_folder := path_strip_filename(get_path_of_running_executable());
  // Lasting copies: temporary storage is reset after every batch.
  out_folder := sprint("%ocr-out", exe_folder);
  i := 2;
  while i < args.count {
    arg := args[i];
    value := ifx i + 1 < args.count then args[i + 1] else "";
    if arg == {
      case "-models";  for split(value, ",")  array_add(*model_names, trim(it)); i += 2;
      case "-cpu";     cpu_only = true; i += 1;
      case "-gpu-only";  gpu_only = true; i += 1;
      case "-workers";   options.gpu_workers = cast(s32) max(1, string_to_int(value)); i += 2;
      case "-rec-batch"; rec_batch = cast(s32) max(1, string_to_int(value)); i += 2;
      case "-warmup";    warmup = max(0, string_to_int(value)); i += 2;
      case "-batch";   batch = max(1, string_to_int(value)); i += 2;
      case "-limit";   limit = max(1, string_to_int(value)); i += 2;
      case "-out";     out_folder = value; i += 2;
      case;            print("Unknown option %.\n%", arg, USAGE); exit(2);
    }
  }
  models_folder := sprint("%ocr-models", exe_folder);
  if !model_names.count {
    models, error := ocr_list_models(models_folder);
    if error { print("%\n", error); exit(1); }
    for models  array_add(*model_names, it.name);
  }

  // The captures, decoded once: the original image of each, as the camera took it.
  Capture :: struct { name: string; pixels: *u8; width, height: s32; }
  captures: [..] Capture;
  decode_start := current_time_monotonic();
  for file_list(folder) {
    if captures.count >= limit  break;
    if !ends_with_nocase(it, ".mcap")  continue;
    data, ok := read_entire_file(it);
    if !ok  continue;
    m, parsed, error := mcap_parse(cast([] u8) data);
    if !parsed { print("Skipping %: %\n", it, error); continue; }
    w, h, channels: s32;
    pixels := stbi_load_from_memory(m.original.data, cast(s32) m.original.count, *w, *h, *channels, 4);
    if !pixels { print("Skipping %: its image cannot be decoded.\n", it); continue; }
    array_add(*captures, .{copy_string(path_strip_extension(path_filename(it))), pixels, w, h});
    free(data);
  }
  if !captures.count { print("No .mcap files in %.\n", folder); exit(1); }
  quick_sort(captures, (a: Capture, b: Capture) -> s64 { return compare(a.name, b.name); });
  megapixels: float64;
  for captures  megapixels += cast(float64) it.width * it.height / 1_000_000;
  print("% captures (% megapixels) decoded in % s.\n\n", captures.count, FF(megapixels, 1), FF(seconds_since(decode_start), 2));

  images := NewArray(captures.count, Ocr_Image);
  for captures  images[it_index] = .{it.pixels, it.width, it.height, it.width * 4};

  for model_name: model_names {
    for prefer_gpu: bool.[true, false] {
      if prefer_gpu && cpu_only  continue;
      if !prefer_gpu && gpu_only  continue;
      create_start := current_time_monotonic();
      engine, ok, error := ocr_create(models_folder, model_name, prefer_gpu, options);
      if !ok { print("% : %\n\n", model_name, error); break; }
      if rec_batch  engine.rec_batch = rec_batch;
      create_seconds := seconds_since(create_start);
      in_use, explanation := ocr_gpu_status(engine);
      provider := ocr_provider(engine);
      if prefer_gpu && !in_use {
        // Said once; the CPU run that follows is the same thing.
        print("%\n\n", explanation);
        ocr_destroy(engine);
        continue;
      }
      print("== % on % (sessions created in % s, % worker threads, % lines per recognizer run)\n", model_name, provider, FF(create_seconds, 2), engine.workers, engine.rec_batch);

      text_folder := sprint("%/%-%", out_folder, model_name, to_lower_copy(tprint("%", provider),, temp));
      make_directory_if_it_does_not_exist(text_folder, recursive = true);

      // The first image alone: the first run pays for setting up kernels and memory.
      first_start := current_time_monotonic();
      ocr_process(engine, array_view(images, 0, 1),, temp);
      first_seconds := seconds_since(first_start);
      // Untimed passes over everything, so the timed one does not pay for setting up each new shape.
      for 1..warmup {
        for start: 0..images.count - 1 {
          if start % batch  continue;
          ocr_process(engine, array_view(images, start, min(batch, images.count - start)),, temp);
          reset_temporary_storage();
        }
      }
      ocr_reset_stats(engine);

      start := current_time_monotonic();
      lines, characters, failures: s64;
      done := 0;
      while done < images.count {
        count := min(batch, images.count - done);
        results := ocr_process(engine, array_view(images, done, count),, temp);
        for result, n: results {
          name := captures[done + n].name;
          if result.failed { failures += 1; print("  % failed: %\n", name, result.error); continue; }
          builder: String_Builder;
          builder.allocator = temp;
          for result.lines {
            print_to_builder(*builder, "% % % %  %  %\n", it.x, it.y, it.width, it.height, FF(it.confidence, 3), it.text);
            characters += it.text.count;
          }
          lines += result.lines.count;
          write_entire_file(tprint("%/%.txt", text_folder, name), *builder);
        }
        done += count;
        reset_temporary_storage();
      }
      seconds := seconds_since(start);
      stats := ocr_stats(engine);
      print("   first image alone:   % s (warm-up)\n", FF(first_seconds, 3));
      print("   all % images:        % s, % ms per image, % megapixels/s, % failed\n", images.count, FF(seconds, 2), FF(1000 * seconds / images.count, 1), FF(megapixels / seconds, 2), failures);
      print("   found:               % lines, % bytes of text\n", lines, characters);
      print("   stage time (summed over threads): detection % s, orientation % s, recognition % s\n", FF(stats.detect_us / 1_000_000.0, 2), FF(stats.orientation_us / 1_000_000.0, 2), FF(stats.recognize_us / 1_000_000.0, 2));
      print("   lines read better turned: %\n", stats.turned);
      print("   text in %\n\n", text_folder);
      ocr_destroy(engine);
    }
  }
}

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

S64_MAX_VALUE :: 0x7FFF_FFFF_FFFF_FFFF;

#import "Basic";
#import "String";
#import "File";
#import "File_Utilities";
#import "Sort";
#import "System";
#import "stb_image";
#import "Capture_File";
#import "OCR";