// Recognition: the text of each detected line. Each box is cropped from the image at full size,
// scaled to rec_height keeping its shape, and normalized to -1..1 in BGR order, as PaddleOCR's
// recognizers expect. Lines of similar width are batched, padded to the widest in the batch.
// The output is a probability per character class per step; greedy CTC decoding takes the best
// class at each step, merges repeats and drops blanks.

Line_Job :: struct {
  image: *Ocr_Image;
  box: Detected_Box;
  width: s32;            // after scaling to rec_height
  // Filled in by recognition:
  text: string;          // default allocator
  confidence: float;
}

// Recognizes jobs (one batch) in place.
ocr_recognize_batch :: (engine: *Ocr_Engine, jobs: [] *Line_Job) -> ok: bool, error: string {
  if !jobs.count  return true, "";
  height := engine.model.rec_height;
  width: s32 = 0;
  for jobs  width = max(width, it.width);
  width = max(width, height);   // the network needs a few steps even for one character

  plane := cast(s64) height * width;
  input := NewArray(plane * 3 * jobs.count, float, initialized = true,, context.default_allocator);   // 0 is the padding
  defer array_free(input,, context.default_allocator);

  for job, n: jobs {
    base := input.data + plane * 3 * n;
    box := job.box;
    bw := box.x1 - box.x0;
    bh := box.y1 - box.y0;
    image := job.image;
    // Bilinear: the crops are mostly enlarged (small UI text up to 48 px), where nearest would be blocky.
    x_scale := cast(float) bw / job.width;
    y_scale := cast(float) bh / height;
    for y: 0..height - 1 {
      sy := (y + 0.5) * y_scale - 0.5 + box.y0;
      y0 := clamp(cast(s32) floor(sy), 0, image.height - 1);
      y1 := min(y0 + 1, image.height - 1);
      fy := clamp(sy - y0, 0, 1);
      row0 := image.pixels + cast(s64) y0 * image.stride;
      row1 := image.pixels + cast(s64) y1 * image.stride;
      for x: 0..job.width - 1 {
        sx := (x + 0.5) * x_scale - 0.5 + box.x0;
        x0 := clamp(cast(s32) floor(sx), 0, image.width - 1);
        x1 := min(x0 + 1, image.width - 1);
        fx := clamp(sx - x0, 0, 1);
        i := cast(s64) y * width + x;
        for c: 0..2 {
          channel := 2 - c;   // B, G, R from RGBA
          top    := row0[x0 * 4 + channel] * (1 - fx) + row0[x1 * 4 + channel] * fx;
          bottom := row1[x0 * 4 + channel] * (1 - fx) + row1[x1 * 4 + channel] * fx;
          value  := top * (1 - fy) + bottom * fy;
          base[i + plane * c] = (value / 255.0 - 0.5) / 0.5;
        }
      }
    }
  }

  shape := s64.[jobs.count, 3, height, width];
  out, ok, error := ocr_run(*engine.rec, engine.memory_info, input, shape);
  if !ok  return false, error;
  defer ocr_release_output(*out);
  if out.dim_count != 3 || out.dims[0] != jobs.count  return false, tprint("The recognizer gave an output of shape %.", array_view(out.dims, 0, out.dim_count));
  steps   := out.dims[1];
  classes := out.dims[2];
  if classes != engine.dictionary.count + 2 {
    return false, tprint("The recognizer has % classes, but its dictionary % characters (plus blank and space, % expected).", classes, engine.dictionary.count, engine.dictionary.count + 2);
  }

  for job, n: jobs {
    builder: String_Builder;
    builder.allocator = temp;
    previous := -1;
    probability_sum: float64;
    characters := 0;
    for t: 0..steps - 1 {
      row := out.data + (n * steps + t) * classes;
      best := 0;
      best_p := row[0];
      for c: 1..classes - 1  if row[c] > best_p { best = c; best_p = row[c]; }
      if best != 0 && best != previous {
        // Class 0 is the CTC blank, 1..count the dictionary, and the last one a space.
        if best <= engine.dictionary.count  append(*builder, engine.dictionary[best - 1]);
        else                               append(*builder, " ");
        probability_sum += best_p;
        characters += 1;
      }
      previous = best;
    }
    job.text = copy_string(trim(builder_to_string(*builder,, temp)),, context.default_allocator);
    job.confidence = ifx characters then cast(float) (probability_sum / characters) else 0;
  }
  return true, "";
}

// The width a box is recognized at: rec_height high, keeping its shape.
ocr_line_width :: (box: Detected_Box, rec_height: s32) -> s32 {
  bw := box.x1 - box.x0;
  bh := max(box.y1 - box.y0, 1);
  return clamp(cast(s32) (cast(float) bw * rec_height / bh + 0.5), 1, REC_MAX_WIDTH);
}

REC_MAX_WIDTH :: 4096;   // a line longer than this, at rec_height, is squeezed to fit