// OCR: the text in images, with where each line is. A batch of RGBA images goes in; each comes back
// as its text lines, with a box and a confidence for each.
//
//   engine, ok, error := ocr_create(models_folder, "ppocr-v5-mobile");
//   results := ocr_process(engine, images);
//
// The models are listed in models_folder/models.config (fetch_models.jai downloads them) and run
// with onnxruntime (modules/onnxruntime), on an NVIDIA GPU through CUDA when it can, else on the
// CPU. ocr_gpu_status says why the GPU is not in use, and what to install so it can be.
//
// Version 1 runs PP-OCRv5's two-stage pipeline: detection (detect.jai), then recognition
// (recognize.jai). Text is expected horizontal, as in screenshots.

Ocr_Provider :: enum { CPU; CUDA; }

// Tightly packed or not: stride is the bytes from one row to the next.
Ocr_Image :: struct {
  pixels: *u8;   // RGBA
  width, height, stride: s32;
}

Ocr_Line :: struct {
  x, y, width, height: s32;   // in the image
  text: string;
  confidence: float;          // 0 to 1: the mean probability of its characters
}

Ocr_Result :: struct {
  lines: [] Ocr_Line;         // top to bottom, then left to right
  failed: bool;
  error: string;
}

Ocr_Engine :: struct {
  model: Ocr_Model_Info;
  provider: Ocr_Provider;
  gpu_explanation: string;    // why the GPU is not in use; "" when it is
  env: *OrtEnv;
  memory_info: *OrtMemoryInfo;
  det, rec: Ocr_Session;
  dictionary: [] string;      // class n is dictionary[n - 1]
  workers: s32;               // threads ocr_process uses
  stats: Ocr_Stats;           // since creation, or the last ocr_reset_stats
}

// Time spent in each stage, summed over the worker threads (so it can exceed the elapsed time).
Ocr_Stats :: struct {
  images: s64;
  lines: s64;
  detect_us: s64;             // preparing images, the detector, and finding the boxes
  recognize_us: s64;          // cropping lines, the recognizer, and decoding the text
}

ocr_stats :: (engine: *Ocr_Engine) -> Ocr_Stats { return engine.stats; }
ocr_reset_stats :: (engine: *Ocr_Engine) { engine.stats = .{}; }

// An engine for the model named in models_folder/models.config. With prefer_gpu it runs on CUDA when
// every library that needs is present; otherwise, or when CUDA fails to start, on the CPU.
ocr_create :: (models_folder: string, model_name: string, prefer_gpu := true) -> *Ocr_Engine, ok: bool, error: string {
  push_allocator(context.default_allocator);

  loaded, load_error := ort_load();
  if !loaded  return null, false, load_error;

  models, models_error := ocr_list_models(models_folder);
  if models_error  return null, false, models_error;
  model: *Ocr_Model_Info;
  for * models  if it.name == model_name { model = it; break; }
  if !model  return null, false, tprint("% has no [model %].", OCR_MODELS_CONFIG, model_name);

  engine := New(Ocr_Engine);
  engine.model = model.*;

  dictionary_text, dictionary_ok := read_entire_file(tprint("%/%", models_folder, model.dictionary), log_errors = false);
  if !dictionary_ok { ocr_destroy(engine); return null, false, tprint("Cannot read the dictionary %/%.", models_folder, model.dictionary); }
  entries := split(dictionary_text, "\n");
  for * entries  if it.count && it.*[it.count - 1] == #char "\r"  it.count -= 1;
  while entries.count && !entries[entries.count - 1].count  entries.count -= 1;   // the final newline
  engine.dictionary = entries;

  api := ort.api;
  ok, error := ort_check(api.CreateEnv(.ORT_LOGGING_LEVEL_WARNING, "ocr", *engine.env));
  if !ok { ocr_destroy(engine); return null, false, error; }
  ort_check(api.CreateCpuMemoryInfo(.OrtArenaAllocator, .OrtMemTypeDefault, *engine.memory_info));

  cores := cast(s32) max(get_number_of_processors(.ALL_LOGICAL), 1);   // efficiency cores included
  detector_path   := tprint("%/%", models_folder, model.detector);
  recognizer_path := tprint("%/%", models_folder, model.recognizer);

  use_gpu := false;
  if prefer_gpu {
    available, explanation := ocr_cuda_check();
    use_gpu = available;
    if !available  engine.gpu_explanation = copy_string(explanation);
  } else {
    engine.gpu_explanation = copy_string("OCR runs on the CPU, as it was asked to.");
  }

  if use_gpu {
    // A few CPU threads prepare images and decode results while the GPU runs the networks.
    engine.workers = 4;
    det_ok, rec_ok: bool;
    det_error, rec_error: string;
    engine.det, det_ok, det_error = ocr_create_session(engine.env, detector_path,   true, 2);
    if det_ok  engine.rec, rec_ok, rec_error = ocr_create_session(engine.env, recognizer_path, true, 2);
    if det_ok && rec_ok {
      engine.provider = .CUDA;
    } else {
      ocr_release_session(*engine.det);
      ocr_release_session(*engine.rec);
      engine.gpu_explanation = sprint("OCR runs on the CPU: the CUDA libraries were found, but onnxruntime could not start on the GPU:\n%", ifx det_ok then rec_error else det_error);
      use_gpu = false;
    }
  }
  if !use_gpu {
    // Several images at once, each network run spread over the rest of the cores.
    engine.workers = clamp(cores / 4, 1, 8);
    threads := max(cores / engine.workers, 1);
    det_ok, rec_ok: bool;
    det_error, rec_error: string;
    engine.det, det_ok, det_error = ocr_create_session(engine.env, detector_path,   false, threads);
    if !det_ok { ocr_destroy(engine); return null, false, det_error; }
    engine.rec, rec_ok, rec_error = ocr_create_session(engine.env, recognizer_path, false, threads);
    if !rec_ok { ocr_destroy(engine); return null, false, rec_error; }
    engine.provider = .CPU;
  }
  return engine, true, "";
}

ocr_provider :: (engine: *Ocr_Engine) -> Ocr_Provider {
  return engine.provider;
}

// Whether the GPU is in use, and when not, why, and what to install so it can be.
ocr_gpu_status :: (engine: *Ocr_Engine) -> in_use: bool, explanation: string {
  return engine.provider == .CUDA, engine.gpu_explanation;
}

ocr_destroy :: (engine: *Ocr_Engine) {
  if !engine  return;
  ocr_release_session(*engine.det);
  ocr_release_session(*engine.rec);
  if engine.memory_info  ort.api.ReleaseMemoryInfo(engine.memory_info);
  if engine.env          ort.api.ReleaseEnv(engine.env);
  free(engine.gpu_explanation,, context.default_allocator);
  free(engine,, context.default_allocator);   // the dictionary and model strings are kept: small, and shared with ocr_list_models' caller
}

// The text in each image. Results (lines and their text) come from allocator. An image that fails
// is marked failed with the reason; the others are unaffected.
ocr_process :: (engine: *Ocr_Engine, images: [] Ocr_Image, allocator := context.allocator) -> [] Ocr_Result {
  results := NewArray(images.count, Ocr_Result,, allocator);
  if !images.count  return results;

  // Detection: images shared out among the workers.
  per_image := NewArray(images.count, Image_Work,, context.default_allocator);
  defer {
    for * per_image  array_free(it.boxes);
    array_free(per_image,, context.default_allocator);
  }
  for * per_image  it.image = *images[it_index];
  run_on_workers(engine, per_image.count, (engine: *Ocr_Engine, index: s64, data: *void) {
    work := cast(*Image_Work) data + index;
    start := current_time_monotonic();
    work.boxes, work.ok, work.error = ocr_detect(engine, work.image.*);
    atomic_add(*engine.stats.detect_us, to_microseconds(current_time_monotonic() - start));
    if work.error  work.error = copy_string(work.error,, context.default_allocator);
  }, per_image.data);

  // Recognition: every line of every image, batched by width so little of each batch is padding.
  jobs: [..] Line_Job;
  jobs.allocator = context.default_allocator;
  defer {
    for jobs  free(it.text,, context.default_allocator);
    array_free(jobs);
  }
  for work: per_image {
    if !work.ok  continue;
    for work.boxes  array_add(*jobs, .{image = work.image, box = it, width = ocr_line_width(it, engine.model.rec_height)});
  }
  order := NewArray(jobs.count, *Line_Job,, context.default_allocator);
  defer array_free(order,, context.default_allocator);
  for * jobs  order[it_index] = it;
  quick_sort(order, (a: *Line_Job, b: *Line_Job) -> s64 { return a.width - b.width; });

  batch_size := max(engine.model.rec_batch, 1);
  batch_count := (order.count + batch_size - 1) / batch_size;
  Batch_Context :: struct { order: [] *Line_Job; size: s64; failed: bool; error: string; }
  batches := Batch_Context.{order = order, size = batch_size};
  run_on_workers(engine, batch_count, (engine: *Ocr_Engine, index: s64, data: *void) {
    b := cast(*Batch_Context) data;
    first := index * b.size;
    count := min(b.size, b.order.count - first);
    start := current_time_monotonic();
    ok, error := ocr_recognize_batch(engine, array_view(b.order, first, count));
    atomic_add(*engine.stats.recognize_us, to_microseconds(current_time_monotonic() - start));
    if !ok { b.failed = true; b.error = copy_string(error,, context.default_allocator); }   // the last error wins; any is enough to say why
  }, *batches);

  // Each image's lines, in reading order.
  for work, image_index: per_image {
    result := *results[image_index];
    if !work.ok {
      result.failed = true;
      result.error = copy_string(work.error,, allocator);
      free(work.error,, context.default_allocator);
      continue;
    }
    if batches.failed {
      result.failed = true;
      result.error = copy_string(batches.error,, allocator);
      continue;
    }
    lines: [..] Ocr_Line;
    lines.allocator = allocator;
    for job: jobs {
      if job.image != work.image || !job.text.count  continue;
      array_add(*lines, .{x = job.box.x0, y = job.box.y0, width = job.box.x1 - job.box.x0, height = job.box.y1 - job.box.y0,
                          text = copy_string(job.text,, allocator), confidence = job.confidence});
    }
    // Rows by the middle of each line, then left to right within a row.
    quick_sort(lines, (a: Ocr_Line, b: Ocr_Line) -> s64 {
      ay, by := a.y + a.height / 2, b.y + b.height / 2;
      if abs(ay - by) > min(a.height, b.height) / 2  return ifx ay < by then -1 else 1;
      return a.x - b.x;
    });
    result.lines = lines;
    engine.stats.lines += lines.count;
  }
  engine.stats.images += images.count;
  free(batches.error,, context.default_allocator);
  return results;
}

#scope_module

Image_Work :: struct {
  image: *Ocr_Image;
  boxes: [..] Detected_Box;
  ok: bool;
  error: string;
}

Work_Proc :: #type (engine: *Ocr_Engine, index: s64, data: *void);

// Calls proc for each index below count, on engine.workers threads at once.
run_on_workers :: (engine: *Ocr_Engine, count: s64, proc: Work_Proc, data: *void) {
  if count <= 0  return;
  Shared :: struct { engine: *Ocr_Engine; count: s64; next: s64; proc: Work_Proc; data: *void; }
  shared := Shared.{engine = engine, count = count, proc = proc, data = data};

  worker_proc :: (thread: *Thread) -> s64 {
    // The context comes from the calling thread, allocator included, and that may be its temporary
    // allocator, which belongs to that thread alone. Workers allocate from the heap, and use their
    // own temporary storage.
    context.allocator = context.default_allocator;
    shared := cast(*Shared) thread.data;
    while true {
      index := atomic_add(*shared.next, 1);
      if index >= shared.count  break;
      shared.proc(shared.engine, index, shared.data);
      reset_temporary_storage();
    }
    return 0;
  }

  // thread_init allocates the thread's temporary storage from the context allocator and tells it
  // the thread is starting; the caller's may be a temporary allocator, which refuses that.
  push_allocator(context.default_allocator);
  thread_count := min(cast(s64) engine.workers, count);
  threads := NewArray(thread_count, Thread,, context.default_allocator);
  defer array_free(threads,, context.default_allocator);
  started := 0;
  for * threads {
    if !thread_init(it, worker_proc, temporary_storage_size = 1024 * 1024)  break;
    thread_start(it, *shared);   // sets thread.data
    started += 1;
  }
  if !started {   // no threads: do it here
    here: Thread;
    here.data = *shared;
    worker_proc(*here);
  }
  for 0..started - 1 {
    thread_is_done(*threads[it], -1);
    thread_deinit(*threads[it]);
  }
}

#load "models.jai";
#load "detect.jai";
#load "recognize.jai";
#load "cuda_check.jai";

#import "Basic";
#import "String";
#import "File";
#import "File_Utilities";
#import "Math";
#import "Sort";
#import "System";
#import "Thread";
#import "Atomics";
#import "onnxruntime";
#if OS == .WINDOWS {
  #import "Windows";
  #import "Windows_Utf8";
} else {
  #import "POSIX";
}