// 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.
//
// It runs PP-OCRv5's pipeline: detection (detect.jai) finds each line as a rectangle at any angle,
// then recognition (recognize.jai) reads it level. When the model lists an orientation classifier,
// lines that look upside down are read both ways; lines read downwards or upwards are read too.

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: the bounds of corners
  text: string;
  confidence: float;          // 0 to 1: the mean probability of its characters
  corners: [4] Vector2;       // in the image's pixels: the text's top left, top right, bottom right, bottom left
}

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;
  orientation: Ocr_Session;   // the text line orientation classifier; none when its session is null
  orientation_width, orientation_height: s32;   // its input size
  dictionary: [] string;      // class n is dictionary[n - 1]
  workers: s32;               // threads ocr_process uses
  rec_batch: s32;             // text lines per recognizer run: the model's rec_batch_gpu or rec_batch_cpu
  options: Ocr_Options;
  stats: Ocr_Stats;           // since creation, or the last ocr_reset_stats
  pool: Ocr_Pool;             // the worker threads, for the engine's life (run_on_workers)
  budget: Ocr_Budget;         // on CUDA: how many runs fit in GPU memory at once (budget.jai)

  // When set, and true, ocr_process stops starting work: the batch of lines or image each worker is
  // on is finished, and then every image fails as "cancelled" (none is part read). Set by the caller.
  cancel: *bool;
}

OCR_CANCELLED :: "cancelled";

ocr_cancelled :: (engine: *Ocr_Engine) -> bool {
  return engine.cancel && engine.cancel.*;
}

// 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
  orientation_us: s64;        // cropping lines and the orientation classifier
  turned: s64;                // lines that read better the other way up, or level as found
  budget_waits: s64;          // detector and recognizer runs that waited for GPU memory (budget.jai)
}

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

// How an engine runs on the GPU, beyond its model's models.config entry. Nearly every image and
// batch is a different size. onnxruntime's arena keeps each block it takes from CUDA and never joins
// or frees them, so inputs that grow one run after another leave a block of every size behind: the
// runs go largest first (ocr_process) so later ones fit in what the first took.
Ocr_Options :: struct {
  gpu_workers: s32 = 4;           // threads preparing images and decoding results while the GPU runs the networks
  cuda_arena_extend_strategy := "kSameAsRequested";  // how much the arena takes when it runs short: what is asked; "kNextPowerOfTwo" takes twice as much each time, up to twice what the runs need
  cudnn_max_workspace := false;   // let cuDNN take as much working memory as it finds best (onnxruntime's own default)
  rec_width_step: s32 = 64;       // a recognizer batch's width is rounded up to a multiple of this, so shapes recur
}

// 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, options := Ocr_Options.{}) -> *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.*;
  engine.options = options;

  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.log_level, "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

  use_gpu := false;
  if prefer_gpu {
    available, explanation := ort_cuda_check("OCR");   // modules/onnxruntime/cuda.jai
    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 = max(options.gpu_workers, 1);
    ort_gpu_lock();   // one CUDA engine made or released at a time (modules/onnxruntime/session.jai)
    sessions_ok, sessions_error := ocr_create_sessions(engine, models_folder, true, 2);
    if sessions_ok  sessions_ok, sessions_error = ocr_warm_up(engine);
    if !sessions_ok  ocr_release_sessions(engine);
    if sessions_ok {
      engine.provider = .CUDA;   // before calibrating: ocr_shrink_arena only works on CUDA
      ocr_budget_calibrate(engine);
    }
    ort_gpu_unlock();
    if sessions_ok {
      engine.provider = .CUDA;
      engine.rec_batch = max(engine.model.rec_batch_gpu, 1);
    } else {
      engine.gpu_explanation = sprint("OCR runs on the CPU: the CUDA libraries were found, but onnxruntime could not start on the GPU:\n%", sessions_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);
    sessions_ok, sessions_error := ocr_create_sessions(engine, models_folder, false, threads);
    if !sessions_ok { ocr_destroy(engine); return null, false, sessions_error; }
    engine.provider = .CPU;
    engine.rec_batch = max(engine.model.rec_batch_cpu, 1);
  }
  ocr_pool_start(engine);
  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_pool_stop(engine);
  ocr_budget_destroy(engine);
  if engine.provider == .CUDA  ort_gpu_lock();
  ocr_release_sessions(engine);
  if engine.provider == .CUDA  ort_gpu_unlock();
  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;
  ocr_budget_update(engine);   // budget.jai: the GPU memory this call's runs may share

  // 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];
  // The largest first: the detector's arena then holds a block that fits every later image
  // (Ocr_Options). Smallest first, it took a new one for each larger image and kept them all.
  detect_order := NewArray(per_image.count, *Image_Work,, context.default_allocator);
  defer array_free(detect_order,, context.default_allocator);
  for * per_image  detect_order[it_index] = it;
  quick_sort(detect_order, (a: *Image_Work, b: *Image_Work) -> s64 {
    return cast(s64) b.image.width * b.image.height - cast(s64) a.image.width * a.image.height;
  });
  run_on_workers(engine, detect_order.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);
  }, detect_order.data);
  ocr_shrink_arena(engine, *engine.det);   // before recognition, so the two stages' blocks are not held at once
  if ocr_cancelled(engine)  return ocr_all_cancelled(results, allocator);

  // 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);
  }
  rec_height := engine.model.rec_height;
  turned: [..] bool;   // per first job: whether detection turned its box a quarter
  turned.allocator = context.default_allocator;
  defer array_free(turned);
  for work: per_image {
    if !work.ok  continue;
    for work.boxes {
      array_add(*jobs, .{image = work.image, corners = it.corners, width = ocr_line_width(it.corners, rec_height)});
      array_add(*turned, it.turned);
    }
  }

  // Other ways to read a line, recognized beside it, the more confident reading kept: upside down,
  // where the orientation classifier finds it more likely than not; and level as it was found, for a
  // box detection turned a quarter, since a narrow word such as "1" is as tall as a line read downwards.
  first_count := jobs.count;
  upside_down := NewArray(first_count, bool,, context.default_allocator);
  defer array_free(upside_down,, context.default_allocator);
  Orientation_Context :: struct { jobs: [] Line_Job; upside_down: [] bool; size: s64; failed: bool; error: string; }
  classify := Orientation_Context.{jobs = jobs, upside_down = upside_down, size = max(engine.rec_batch, 1)};
  if engine.orientation.session && first_count {
    run_on_workers(engine, (first_count + classify.size - 1) / classify.size, (engine: *Ocr_Engine, index: s64, data: *void) {
      c := cast(*Orientation_Context) data;
      first := index * c.size;
      count := min(c.size, c.jobs.count - first);
      batch := NewArray(count, *Line_Job,, temp);
      for * batch  it.* = *c.jobs[first + it_index];
      start := current_time_monotonic();
      ok, error := ocr_classify_batch(engine, batch, array_view(c.upside_down, first, count));
      atomic_add(*engine.stats.orientation_us, to_microseconds(current_time_monotonic() - start));
      if !ok { c.failed = true; c.error = copy_string(error,, context.default_allocator); }   // the last error wins
    }, *classify);
  }
  defer free(classify.error,, context.default_allocator);
  for n: 0..first_count - 1 {
    job := jobs[n];
    if upside_down[n]  array_add(*jobs, .{image = job.image, corners = ocr_corners_turned(job.corners, 2), width = job.width, primary = n});
    if turned[n] {
      level := ocr_corners_turned(job.corners, 3);
      array_add(*jobs, .{image = job.image, corners = level, width = ocr_line_width(level, rec_height), primary = n});
    }
  }

  order := NewArray(jobs.count, *Line_Job,, context.default_allocator);
  defer array_free(order,, context.default_allocator);
  for * jobs  order[it_index] = it;
  // Widest first, so each batch fits in the arena block the first one took (Ocr_Options). Narrowest
  // first, every batch was wider than the last and took a new block: 5.7 GB of them for the mobile
  // model on gpu-bench's 30 images, where the widest batch needs 0.9 GB.
  quick_sort(order, (a: *Line_Job, b: *Line_Job) -> s64 { return b.width - a.width; });

  batch_size := max(engine.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);
  if ocr_cancelled(engine) {
    free(batches.error,, context.default_allocator);
    return ocr_all_cancelled(results, allocator);
  }
  if classify.failed && !batches.failed {
    batches.failed = true;
    batches.error = copy_string(classify.error,, context.default_allocator);
  }

  // The more confident of each line's readings takes the first job's place.
  for * other: array_view(jobs, first_count, jobs.count - first_count) {
    first := *jobs[other.primary];
    if other.confidence > first.confidence {
      swap(*first.text, *other.text);
      first.corners, first.confidence = other.corners, other.confidence;
      engine.stats.turned += 1;
    }
  }

  ocr_shrink_arena(engine, *engine.rec);   // models.jai: what this call's runs took, given back

  // 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: array_view(jobs, 0, first_count) {
      if job.image != work.image || !job.text.count  continue;
      x0, y0, x1, y1 := ocr_corner_bounds(job.corners, job.image.width, job.image.height);
      array_add(*lines, .{x = x0, y = y0, width = x1 - x0, height = y1 - y0, text = copy_string(job.text,, allocator),
                          confidence = job.confidence, corners = job.corners});
    }
    // 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

ocr_all_cancelled :: (results: [] Ocr_Result, allocator: Allocator) -> [] Ocr_Result {
  for * results {
    it.failed = true;
    it.error = copy_string(OCR_CANCELLED,, allocator);
  }
  return results;
}

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.
// The engine's worker threads: started with it and kept until ocr_destroy, never one set per call.
// onnxruntime (1.26 to 1.30 at least) sets up a CUDA stream with cuBLAS and cuDNN handles for every
// OS thread that runs a session, about 10 MB each, kept after the thread ends (microsoft/onnxruntime
// #33037). Workers made afresh for every stage of every call, as they were, gave the OCR sessions
// thousands of threads in a scan, and starting them cost time: 6% on gpu-bench's mobile OCR.
Ocr_Pool :: struct {
  threads: [] Thread;
  started: s64;
  mutex: Mutex;              // one run_on_workers at a time
  wake: Semaphore;           // signalled once per worker for each job, and for quitting
  done: Semaphore;           // each worker signals it when it has finished its part of a job
  quit: bool;
  // The job: items 0..count-1, taken in turn.
  count, next: s64;
  proc: Work_Proc;
  data: *void;
}

ocr_pool_start :: (engine: *Ocr_Engine) {
  pool := *engine.pool;
  init(*pool.mutex);
  init(*pool.wake);
  init(*pool.done);
  // thread_init allocates the thread's temporary storage from the context allocator.
  push_allocator(context.default_allocator);
  pool.threads = NewArray(max(engine.workers, 1), Thread);
  for * pool.threads {
    if !thread_init(it, ocr_pool_worker, temporary_storage_size = 1024 * 1024)  break;
    thread_start(it, engine);   // sets thread.data
    pool.started += 1;
  }
}

ocr_pool_stop :: (engine: *Ocr_Engine) {
  pool := *engine.pool;
  if !pool.threads.count  return;
  pool.quit = true;
  for 1..pool.started  signal(*pool.wake);
  for 0..pool.started - 1 {
    thread_is_done(*pool.threads[it], -1);
    thread_deinit(*pool.threads[it]);
  }
  array_free(pool.threads,, context.default_allocator);
  destroy(*pool.wake);
  destroy(*pool.done);
  destroy(*pool.mutex);
  pool.* = .{};
}

ocr_pool_worker :: (thread: *Thread) -> s64 {
  // Workers allocate from the heap, and use their own temporary storage.
  context.allocator = context.default_allocator;
  engine := cast(*Ocr_Engine) thread.data;
  pool := *engine.pool;
  while true {
    Default_Allocator.thread_idle();   // parked until the next job: give back what this thread's heap holds free
    wait_for(*pool.wake);
    if pool.quit  break;
    while true {
      if ocr_cancelled(engine)  break;   // what is left fails as cancelled (ocr_process)
      index := atomic_add(*pool.next, 1);
      if index >= pool.count  break;
      pool.proc(engine, index, pool.data);
      reset_temporary_storage();
    }
    signal(*pool.done);
  }
  return 0;
}

run_on_workers :: (engine: *Ocr_Engine, count: s64, proc: Work_Proc, data: *void) {
  if count <= 0  return;
  pool := *engine.pool;
  if !pool.started {   // no threads: do it here
    for 0..count - 1 {
      if ocr_cancelled(engine)  break;
      proc(engine, it, data);
    }
    return;
  }
  lock(*pool.mutex);
  defer unlock(*pool.mutex);
  pool.count, pool.next, pool.proc, pool.data = count, 0, proc, data;
  // Every worker is woken and reports back, even with fewer items than workers: a worker that finds
  // nothing left reports at once, so no wake is left over for the next job.
  for 1..pool.started  signal(*pool.wake);
  for 1..pool.started  wait_for(*pool.done);
}

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

#import "Basic";
#import "String";
#import "File";
#import "File_Utilities";
#import "Math";
#import "Sort";
#import "System";
#import "Thread";
#import "Atomics";
#import "onnxruntime";
Default_Allocator :: #import "Default_Allocator";   // rpmalloc (modules/Default_Allocator), for thread_idle where the workers wait
#if OS == .WINDOWS {
  #import "Windows";
  #import "Windows_Utf8";
} else {
  #import "POSIX";
}