// Detection: where the text lines are. The detector (a DB network) gives each pixel a probability
// of being text. Pixels above det_threshold are grouped into connected regions, and each kept region
// becomes the smallest rectangle around it, at whatever angle, grown by the DB unclip distance, since
// the network is trained to mark a shrunken core of each line.
//
// Images wider than det_max_side are scaled down for detection. Images taller than det_tile_height
// (after scaling) are detected in overlapping strips, so a long scrolling capture keeps its text at
// full size; each box is kept by the strip whose own band holds its centre.

Detected_Box :: struct {
  // In the image's pixels, clockwise: the text's top left, top right, bottom right, bottom left, as
  // far as its shape tells. A box at least TURN_RATIO times as tall as it is wide is taken as a line
  // read downwards, so its first corner is its top right. Recognition may still turn it by 180°.
  corners: [4] Vector2;
  turned: bool;          // taken as a line read downwards: started from its top right
  score: float;
}

TURN_RATIO :: 1.5;

// Screenshots are mostly level: a region keeps a level box unless a turned one is this much smaller.
LEVEL_PREFERENCE :: 1.2;   // a long line tilted by 1° already has a level box ~30% bigger

DET_TILE_OVERLAP : s32 : 128;   // px of each strip shared with the next, after scaling

// The image's text boxes, from the default allocator.
ocr_detect :: (engine: *Ocr_Engine, image: Ocr_Image) -> [..] Detected_Box, ok: bool, error: string {
  boxes: [..] Detected_Box;
  boxes.allocator = context.default_allocator;
  model := *engine.model;
  if image.width <= 0 || image.height <= 0  return boxes, true, "";

  scale := min(1.0, cast(float) model.det_max_side / image.width);
  scaled_width  := max(cast(s32) 1, cast(s32) (image.width  * scale + 0.5));
  scaled_height := max(cast(s32) 1, cast(s32) (image.height * scale + 0.5));
  tile_height   := max(model.det_tile_height, cast(s32) DET_TILE_OVERLAP * 2);

  tile_top: s32 = 0;
  while tile_top < scaled_height {
    tile_bottom := min(tile_top + tile_height, scaled_height);
    // This strip's own band: what is not left to the strips above and below it.
    band_top    := ifx tile_top == 0 then 0 else tile_top + DET_TILE_OVERLAP / 2;
    band_bottom := ifx tile_bottom == scaled_height then scaled_height else tile_bottom - DET_TILE_OVERLAP / 2;

    ok, error := ocr_detect_tile(engine, image, scale, scaled_width, tile_top, tile_bottom, band_top, band_bottom, *boxes);
    if !ok  return boxes, false, error;
    if tile_bottom == scaled_height  break;
    tile_top = tile_bottom - DET_TILE_OVERLAP;
  }
  return boxes, true, "";
}

#scope_module

// Normalization the PP-OCR detectors were trained with: ImageNet mean and deviation, applied to the
// channels in BGR order, as PaddleOCR reads images with OpenCV.
DET_MEAN :: float.[0.485, 0.456, 0.406];
DET_STD  :: float.[0.229, 0.224, 0.225];

ocr_detect_tile :: (engine: *Ocr_Engine, image: Ocr_Image, scale: float, scaled_width: s32, tile_top: s32, tile_bottom: s32, band_top: s32, band_bottom: s32, boxes: *[..] Detected_Box) -> ok: bool, error: string {
  model := *engine.model;
  valid_w := scaled_width;
  valid_h := tile_bottom - tile_top;
  w := align_up_32(valid_w);
  h := align_up_32(valid_h);

  // Planar BGR, normalized; the padding to a multiple of 32 is the normalized zero (the mean colour).
  plane := cast(s64) w * h;
  input := NewArray(plane * 3, float, initialized = true,, context.default_allocator);
  defer array_free(input,, context.default_allocator);
  inverse := 1.0 / scale;
  for y: 0..valid_h - 1 {
    source_y := min(cast(s32) ((y + tile_top + 0.5) * inverse), image.height - 1);
    row := image.pixels + cast(s64) source_y * image.stride;
    for x: 0..valid_w - 1 {
      source_x := min(cast(s32) ((x + 0.5) * inverse), image.width - 1);
      p := row + source_x * 4;
      i := cast(s64) y * w + x;
      input[i]             = (p[2] / 255.0 - DET_MEAN[0]) / DET_STD[0];   // B
      input[i + plane]     = (p[1] / 255.0 - DET_MEAN[1]) / DET_STD[1];   // G
      input[i + plane * 2] = (p[0] / 255.0 - DET_MEAN[2]) / DET_STD[2];   // R
    }
  }

  shape := s64.[1, 3, h, w];
  reserved := ocr_budget_detector(engine, h, w);
  ocr_budget_reserve(engine, reserved);   // budget.jai: waits while the GPU's memory is taken
  out, ok, error := ocr_run(*engine.det, engine.memory_info, input, shape);
  ocr_budget_release(engine, reserved);
  if !ok  return false, error;
  defer ocr_release_output(*out);
  if out.dim_count != 4 || out.dims[2] != h || out.dims[3] != w {
    return false, tprint("The detector gave an output of shape %, not [1, 1, %, %].", array_view(out.dims, 0, out.dim_count), h, w);
  }

  regions := ocr_find_regions(out.data, w, valid_w, valid_h, model.det_threshold);
  defer {
    for regions  array_free(it.outline);
    array_free(regions,, context.default_allocator);
  }
  for regions {
    if it.pixels == 0  continue;
    score := cast(float) (it.probability_sum / it.pixels);
    if score < model.box_threshold  continue;

    rect := ocr_min_area_rect(it.outline);
    rect_w := rect.u_max - rect.u_min;
    rect_h := rect.v_max - rect.v_min;
    if min(rect_w, rect_h) < 3  continue;   // specks

    // Grown as DB's unclip does: by area * ratio / perimeter on every side.
    distance := (rect_w * rect_h * model.unclip_ratio) / (2.0 * (rect_w + rect_h));
    rect.u_min -= distance;  rect.u_max += distance;
    rect.v_min -= distance;  rect.v_max += distance;
    corners := ocr_rect_corners(rect);

    centre_y := tile_top + (corners[0].y + corners[2].y) / 2;
    if centre_y < cast(float) band_top || centre_y >= cast(float) band_bottom  continue;   // the neighbouring strip's

    box: Detected_Box;
    box.score = score;
    for corners {
      box.corners[it_index].x = clamp(it.x * inverse, 0, cast(float) image.width);
      box.corners[it_index].y = clamp((it.y + tile_top) * inverse, 0, cast(float) image.height);
    }
    if rect.u.x == 1 && rect.u.y == 0 {
      // A level box takes in every pixel it touches: the recognizers read these crops best.
      x0, y0, x1, y1 := ocr_corner_bounds(box.corners, image.width, image.height);
      box.corners = .[.{xx x0, xx y0}, .{xx x1, xx y0}, .{xx x1, xx y1}, .{xx x0, xx y1}];
    }
    box.corners, box.turned = ocr_reading_order(box.corners);
    x0, y0, x1, y1 := ocr_corner_bounds(box.corners, image.width, image.height);
    if x1 - x0 >= 2 && y1 - y0 >= 2  array_add(boxes, box);
  }
  return true, "";
}

// The pixels the corners cover, clipped to the image: x0 and y0 inclusive, x1 and y1 exclusive.
ocr_corner_bounds :: (corners: [4] Vector2, image_width: s32, image_height: s32) -> x0: s32, y0: s32, x1: s32, y1: s32 {
  low  := corners[0];
  high := corners[0];
  for corners {
    low.x  = min(low.x,  it.x);  low.y  = min(low.y,  it.y);
    high.x = max(high.x, it.x);  high.y = max(high.y, it.y);
  }
  return clamp(cast(s32) floor(low.x), 0, image_width),  clamp(cast(s32) floor(low.y), 0, image_height),
         clamp(cast(s32) ceil(high.x), 0, image_width),  clamp(cast(s32) ceil(high.y), 0, image_height);
}

// The same rectangle's corners started from another one, keeping them clockwise: by 1, turned a
// quarter (what was the right side becomes the top), by 2, upside down.
ocr_corners_turned :: (corners: [4] Vector2, quarters: s64) -> [4] Vector2 {
  result: [4] Vector2;
  for 0..3  result[it] = corners[(it + quarters) % 4];
  return result;
}

// A rectangle as ranges along two perpendicular unit axes: u, and v = u turned a quarter clockwise
// on the screen (y down).
Rotated_Rect :: struct {
  u: Vector2;
  u_min, u_max, v_min, v_max: float;
}

ocr_rect_corners :: (rect: Rotated_Rect) -> [4] Vector2 {
  u := rect.u;
  v := Vector2.{-u.y, u.x};
  return .[u * rect.u_min + v * rect.v_min, u * rect.u_max + v * rect.v_min,
           u * rect.u_max + v * rect.v_max, u * rect.u_min + v * rect.v_max];
}

// The corners clockwise from the text's top left, as PaddleOCR orders them: of the two leftmost, the
// higher is the top left; of the two rightmost, the higher is the top right. A box at least
// TURN_RATIO times as tall as it is wide is then started from its top right, so a line read
// downwards comes out level.
ocr_reading_order :: (corners: [4] Vector2) -> [4] Vector2, turned: bool {
  sorted := corners;
  quick_sort(sorted, (a: Vector2, b: Vector2) -> s64 { return ifx a.x < b.x then -1 else ifx a.x > b.x then 1 else 0; });
  top_left, bottom_left   := sorted[0], sorted[1];
  if bottom_left.y < top_left.y   swap(*top_left, *bottom_left);
  top_right, bottom_right := sorted[2], sorted[3];
  if bottom_right.y < top_right.y swap(*top_right, *bottom_right);
  result := Vector2.[top_left, top_right, bottom_right, bottom_left];

  width  := length(result[1] - result[0]);
  height := length(result[3] - result[0]);
  if height >= width * TURN_RATIO  return ocr_corners_turned(result, 1), true;
  return result, false;
}

// The smallest rectangle holding the points, at any angle, found by rotating calipers: one side of
// it lies along an edge of the points' convex hull. A level rectangle is kept unless a turned one is
// smaller by more than LEVEL_PREFERENCE.
ocr_min_area_rect :: (points: [] Vector2) -> Rotated_Rect {
  level: Rotated_Rect;
  level.u = .{1, 0};
  level.u_min, level.v_min = FLOAT32_MAX, FLOAT32_MAX;
  level.u_max, level.v_max = -FLOAT32_MAX, -FLOAT32_MAX;
  for points {
    level.u_min = min(level.u_min, it.x);  level.u_max = max(level.u_max, it.x);
    level.v_min = min(level.v_min, it.y);  level.v_max = max(level.v_max, it.y);
  }
  if points.count < 3  return level;
  level_area := (level.u_max - level.u_min) * (level.v_max - level.v_min);

  hull := ocr_convex_hull(points);
  best := level;
  best_area := level_area;
  for i: 0..hull.count - 1 {
    edge := hull[(i + 1) % hull.count] - hull[i];
    edge_length := length(edge);
    if edge_length < 0.0001  continue;
    rect: Rotated_Rect;
    rect.u = edge / edge_length;
    v := Vector2.{-rect.u.y, rect.u.x};
    rect.u_min, rect.v_min = FLOAT32_MAX, FLOAT32_MAX;
    rect.u_max, rect.v_max = -FLOAT32_MAX, -FLOAT32_MAX;
    for hull {
      a := dot(it, rect.u);
      b := dot(it, v);
      rect.u_min = min(rect.u_min, a);  rect.u_max = max(rect.u_max, a);
      rect.v_min = min(rect.v_min, b);  rect.v_max = max(rect.v_max, b);
    }
    area := (rect.u_max - rect.u_min) * (rect.v_max - rect.v_min);
    if area < best_area { best = rect; best_area = area; }
  }
  if level_area <= best_area * LEVEL_PREFERENCE  return level;
  return best;
}

// The convex hull of the points, counter-clockwise (Andrew's monotone chain). Temp.
ocr_convex_hull :: (points: [] Vector2) -> [] Vector2 {
  sorted := array_copy(points,, temp);
  quick_sort(sorted, (a: Vector2, b: Vector2) -> s64 {
    if a.x != b.x  return ifx a.x < b.x then -1 else 1;
    return ifx a.y < b.y then -1 else ifx a.y > b.y then 1 else 0;
  });
  cross :: (o: Vector2, a: Vector2, b: Vector2) -> float { return (a.x - o.x) * (b.y - o.y) - (a.y - o.y) * (b.x - o.x); }

  hull := NewArray(sorted.count * 2, Vector2,, temp);
  k := 0;
  for sorted {   // lower hull
    while k >= 2 && cross(hull[k - 2], hull[k - 1], it) <= 0  k -= 1;
    hull[k] = it;  k += 1;
  }
  lower := k + 1;
  for < i: 0..sorted.count - 2 {   // upper hull
    p := sorted[i];
    while k >= lower && cross(hull[k - 2], hull[k - 1], p) <= 0  k -= 1;
    hull[k] = p;  k += 1;
  }
  hull.count = max(k - 1, 1);   // the last point is the first again
  return hull;
}

align_up_32 :: (n: s32) -> s32 {
  return (n + 31) & cast(s32) ~31;
}

Region :: struct {
  pixels: s64;
  probability_sum: float64;
  // The corners of the first and last pixel in each of its rows: what its convex hull needs.
  outline: [..] Vector2;      // default allocator
  row: s32 = -1;              // the row being scanned, and its first and last pixel
  row_first, row_last: s32;
}

// Groups the pixels above threshold into 8-connected regions: one pass assigning provisional labels
// and recording which touch (union-find), then one summing each region's probability and outline.
ocr_find_regions :: (probability: *float, stride: s32, width: s32, height: s32, threshold: float) -> [..] Region {
  regions: [..] Region;
  regions.allocator = context.default_allocator;

  labels := NewArray(cast(s64) width * height, s32, initialized = true,, context.default_allocator);
  defer array_free(labels,, context.default_allocator);
  parent: [..] s32;
  parent.allocator = context.default_allocator;
  defer array_free(parent);
  array_add(*parent, 0);   // label 0 is background

  find :: (parent: [] s32, label: s32) -> s32 {
    root := label;
    while parent[root] != root  root = parent[root];
    while parent[label] != root { next := parent[label]; parent[label] = root; label = next; }   // path compression
    return root;
  }
  unite :: (parent: [] s32, a: s32, b: s32) {
    ra, rb := find(parent, a), find(parent, b);
    if ra < rb  parent[rb] = ra;
    else if rb < ra  parent[ra] = rb;
  }

  flush_row :: (r: *Region) {
    if r.row < 0  return;
    top, bottom := cast(float) r.row, cast(float) r.row + 1;
    array_add(*r.outline, .{cast(float) r.row_first, top}, .{cast(float) r.row_first, bottom},
                          .{cast(float) r.row_last + 1, top}, .{cast(float) r.row_last + 1, bottom});
    r.row = -1;
  }

  for y: 0..height - 1 {
    for x: 0..width - 1 {
      if probability[cast(s64) y * stride + x] <= threshold  continue;
      i := cast(s64) y * width + x;
      label: s32 = 0;
      // Neighbours already visited: left, and the three above.
      neighbours: [4] s32;
      neighbours[0] = ifx x > 0                       then labels[i - 1]         else 0;
      neighbours[1] = ifx y > 0 && x > 0              then labels[i - width - 1] else 0;
      neighbours[2] = ifx y > 0                       then labels[i - width]     else 0;
      neighbours[3] = ifx y > 0 && x < width - 1      then labels[i - width + 1] else 0;
      for neighbours {
        if !it  continue;
        if !label  label = it;
        else if it != label  unite(parent, label, it);
      }
      if !label {
        label = cast(s32) parent.count;
        array_add(*parent, label);
      }
      labels[i] = label;
    }
  }

  // Each root gets a region.
  region_of := NewArray(parent.count, s32, initialized = true,, context.default_allocator);
  defer array_free(region_of,, context.default_allocator);
  for y: 0..height - 1 {
    for x: 0..width - 1 {
      i := cast(s64) y * width + x;
      if !labels[i]  continue;
      root := find(parent, labels[i]);
      index := region_of[root] - 1;
      if index < 0 {
        region := array_add(*regions);
        region.* = .{};
        region.outline.allocator = context.default_allocator;
        index = cast(s32) regions.count - 1;
        region_of[root] = index + 1;
      }
      r := *regions[index];
      if r.row != y {
        flush_row(r);
        r.row, r.row_first = y, x;
      }
      r.row_last = x;   // rows are scanned left to right
      r.pixels += 1;
      r.probability_sum += probability[cast(s64) y * stride + x];
    }
  }
  for * regions  flush_row(it);
  return regions;
}