// Segmentation in the app (modules/Segment): the segmentation scan, the Segmentation window, and
// -segment. A scan reads only images in the library folders set to segment (app/library_folders.jai),
// never a video, and stores each one's segments in its .mcap: a segs chunk in front, which the
// catalog reads, and a mask chunk at the back (modules/Capture_File/segment_chunks.jai).
//
// The scan runs on the OCR thread (app/ocr_ui.jai), as a task of its own, so it shares the OCR scan's
// progress, pause and cancel, and one thread uses the GPU at a time. The engine is made by the first
// scan, not at startup: it takes some 2.4 GB of GPU memory, and most sessions never segment anything.

SEGMENT_WINDOW     :: "Segmentation";
SEGMENT_MODEL      :: "mask2former-swin-base-coco-panoptic";   // the one segmentation model; its chunks are keyed by this
SEGMENT_SCAN_BATCH :: 8;   // images read and segmented together (the engine runs them 4 at a time on the GPU)

segment_ui: struct {
  // Owned by the OCR thread while ocr_ui.thread_busy, like ocr_ui.engine.
  engine: *Segment_Engine;
  engine_error: string;

  // Main thread.
  window_open: bool;
  scan_requested: bool;          // -segment has started its scan
  logged_step: s64;              // -segment's last progress line, in twentieths
  erase_folder: string;          // the folder whose segments the erase popup asks about; heap
  erase_asked: bool;

  // The window's counts per library folder, worked out again when the catalog changes.
  counted_generation: s64 = -1;
  counts: [..] Segment_Folder_Count;

  overlay_state: Segment_Overlay_State;   // the Capture window's overlay; main thread
}

Segment_Overlay_State :: struct {
  show := true;
  opacity: float = 0.45;
  hidden: [..] string;           // classes not drawn, by name, for every capture; heap
  colours: [..] Segment_Colour;  // classes given their own colour in the legend; kept in SEGMENT_COLOURS_FILE
  overlay: *Segment_Overlay;     // the capture's shown, or being replaced
  entry: *Catalog_Entry;         // the catalog entry it was asked for: a new one is read again
  loading: bool;
  texture: Gpu_Texture;
  dirty: bool;                   // the texture is out of date
}

Segment_Folder_Count :: struct {
  folder: string;                // a view of viewer.folders' string, good until they change
  images: s64;                   // captures that are not videos
  segmented: s64;                // of them, with this model's segments
  catalogued: bool;              // every capture in it has its catalog entry
}

// A class's colour, chosen in the Capture window's legend. Heap.
Segment_Colour :: struct {
  class: string;
  rgb: [3] u8;
}

// The colours chosen, one line each: colour: "person" #e05a47. In the app's data folder, beside the
// label list: a way of looking, so the same in every library.
SEGMENT_COLOURS_FILE :: "segment-colours.config";

segment_colours_path :: () -> string {
  return tprint("%/%", app_program_data_dir(), SEGMENT_COLOURS_FILE);
}

segment_colours_load :: () {
  o := *segment_ui.overlay_state;
  o.colours.allocator = context.default_allocator;
  text, ok := read_entire_file(segment_colours_path(), log_errors = false,, temp);
  if !ok  return;
  for raw: split(text, "\n",, temp) {
    line := trim(trim_right(raw, "\r"));
    if !begins_with(line, "colour:")  continue;   // @version 1, comments, and keys a later build adds
    class, rest, class_ok := mcap_read_quoted(trim(advance(line, 7)),, temp);
    rest = trim(rest);
    if !class_ok || rest.count != 7 || rest[0] != #char "#"  continue;
    value, value_ok := string_to_int(advance(rest, 1), 16, u32);
    if !value_ok  continue;
    segment_colour_set(class, .[cast(u8) (value >> 16), cast(u8) ((value >> 8) & 0xFF), cast(u8) (value & 0xFF)], save = false);
  }
}

segment_colours_save :: () {
  builder: String_Builder;
  builder.allocator = temp;
  append(*builder, "@version 1\n# The colours of segment classes in the Capture window, chosen in its legend. A class not\n# listed has its own colour, from its name.\n");
  for segment_ui.overlay_state.colours {
    print_to_builder(*builder, "colour: % #%\n", mcap_quote(it.class), formatInt((cast(u32) it.rgb[0] << 16) | (cast(u32) it.rgb[1] << 8) | it.rgb[2], base = 16, minimum_digits = 6));
  }
  if !write_entire_file(segment_colours_path(), *builder)  log_error("[Segment] Could not write %.", segment_colours_path());
}

// Gives class its own colour; or, with default, takes it back to the one from its name.
segment_colour_set :: (class: string, rgb: [3] u8, default := false, save := true) {
  o := *segment_ui.overlay_state;
  for * o.colours {
    if it.class != class  continue;
    if default {
      free(it.class,, context.default_allocator);
      array_ordered_remove_by_index(*o.colours, it_index);
    } else {
      it.rgb = rgb;
    }
    if save  segment_colours_save();
    o.dirty = true;
    return;
  }
  if default  return;
  array_add(*o.colours, .{copy_string(class,, context.default_allocator), rgb});
  if save  segment_colours_save();
  o.dirty = true;
}

segment_models_folder :: () -> string {
  return tprint("%segment-models", path_strip_filename(get_path_of_running_executable()));
}

// ---------------------------------------------------------------- startup and shutdown

segment_ui_startup :: () {
  segment_ui.counts.allocator = context.default_allocator;
  segment_ui.overlay_state.hidden.allocator = context.default_allocator;
  segment_colours_load();
  library_folders_load();
}

// After ocr_ui_shutdown has stopped the OCR thread.
segment_ui_shutdown :: () {
  o := *segment_ui.overlay_state;   // the viewer's workers have stopped; the texture goes with the device
  segment_overlay_free(o.overlay);
  segment_overlay_free(viewer.segments_done);
  for o.hidden  free(it,, context.default_allocator);
  array_free(o.hidden);
  for o.colours  free(it.class,, context.default_allocator);
  array_free(o.colours);
  segment_destroy(segment_ui.engine);
  segment_ui.engine = null;
  free(segment_ui.engine_error,, context.default_allocator);
  free(segment_ui.erase_folder,, context.default_allocator);
  array_free(segment_ui.counts);
}

// ---------------------------------------------------------------- the scan

// Starts a scan of the images not yet segmented in the folders set to segment, or of those among only
// (an import's). False, with a toast, when there is nothing to do or the OCR thread is busy.
segment_scan_start :: (only: [] s32 = .[], quiet := false) -> bool {
  if ocr_ui.thread_busy  return false;
  scan := *ocr_ui.scan;
  for scan.items  free(it.path);
  array_reset(*scan.items);
  free(scan.model);
  scan.* = .{};
  scan.items.allocator = context.default_allocator;
  scan.model = copy_string(SEGMENT_MODEL,, context.default_allocator);
  scan.segmenting = true;
  current := ifx segment_ui.engine then segment_classes_text(segment_ui.engine) else "";   // the OCR thread is idle: the engine is ours

  indices: [..] s32;
  indices.allocator = temp;
  if only.count  array_add(*indices, ..only);
  else           for viewer.captures  array_add(*indices, cast(s32) it_index);
  for indices {
    capture := *viewer.captures[it];
    if capture.removed || !library_folder_segments(capture.folder)  continue;
    entry := capture.catalog;
    if !entry || entry.media == .VIDEO  continue;   // images only; the catalog says which
    // Done already, unless with another class list than the engine's (when it is up to say).
    if entry.segment_classes && (!current || entry.segment_classes == current)  continue;
    array_add(*scan.items, .{index = it, path = copy_string(capture.path,, context.default_allocator)});
  }
  scan.total = scan.items.count;
  if !scan.total {
    if !quiet  toast_show(.Success, "Every image in the folders set to segment has been segmented");
    return false;
  }
  scan.running = true;
  scan.start = current_time_monotonic();
  return ocr_ui_start(.{kind = .SEGMENT});
}

// On the OCR thread: the SEGMENT task. Makes the engine the first time, then segments the scan's items
// a batch at a time, storing each.
segment_ui_scan_task :: () {
  scan := *ocr_ui.scan;
  defer atomic_swap(*scan.running, false);
  if !segment_ui.engine {
    start := current_time_monotonic();
    engine, ok, error := segment_create(segment_models_folder(), SEGMENT_MODEL);
    free(segment_ui.engine_error);
    segment_ui.engine_error = ifx ok then "" else copy_string(error);
    if !ok {
      log_error("[Segment] Could not start %: %", SEGMENT_MODEL, error);
      atomic_swap(*scan.failed, scan.total);
      atomic_swap(*scan.done, scan.total);
      return;
    }
    segment_ui.engine = engine;
    log("[Segment] % on % in % ms.%", SEGMENT_MODEL, engine.provider, to_milliseconds(current_time_monotonic() - start), ifx engine.gpu_explanation then tprint(" %", engine.gpu_explanation) else "");
  }
  engine := segment_ui.engine;
  // Cancelling (or closing the app) stops it between batches of the engine's, rather than waiting the
  // scan's batch out. What it had not segmented is left as it was.
  engine.cancel = *scan.cancel;
  defer engine.cancel = null;

  position := 0;
  while position < scan.items.count {
    while scan.paused && !scan.cancel  sleep_milliseconds(100);
    if scan.cancel  break;
    count := min(SEGMENT_SCAN_BATCH, scan.items.count - position);
    errors := segment_read_and_store(engine, array_view(scan.items, position, count));
    if scan.cancel {   // cut short: what it stored stays stored and is counted; the rest is left, not failed
      for errors {
        if !it  atomic_add(*scan.done, 1);
        free(it);
      }
      break;
    }
    for errors  if it { atomic_add(*scan.failed, 1); log_error("[Segment scan] % %", scan.items[position + it_index].path, it); free(it); }
    atomic_add(*scan.done, count);
    position += count;
    reset_temporary_storage();
  }
  scan.elapsed_seconds = to_float64_seconds(current_time_monotonic() - scan.start) - scan.paused_seconds;
  log("[Segment scan] % of % images segmented with % in % (% failed)%.", scan.done, scan.total, SEGMENT_MODEL,
      ocr_duration_text(scan.elapsed_seconds), scan.failed, ifx scan.cancel then ", cancelled" else "");
}

// On the OCR thread: segments the items' images and stores each one's segments, handing its new
// catalog entry to the viewer. An error message (heap) for each that failed, "" for the others. A
// video is skipped, not failed: the catalog that chose it may have been out of date.
segment_read_and_store :: (engine: *Segment_Engine, items: [] Scan_Item) -> [] string {
  errors := NewArray(items.count, string);
  Prepared :: struct { image: Image; which: string; ok: bool; }
  prepared := NewArray(items.count, Prepared,, temp);
  defer for * prepared  if it.ok  image_free(*it.image);
  inputs: [..] Segment_Image;
  inputs.allocator = temp;
  input_of: [..] s64;   // which item each input is
  input_of.allocator = temp;
  for item, n: items {
    m, bytes, which, ok, error := mcap_read_shown_image(item.path,, temp);
    if !ok { errors[n] = copy_string(error); continue; }
    if contains(m.text, "VideoCodec:")  continue;   // a video (app/import.jai): images only
    image, decoded, decode_error := image_decode(bytes);   // turned by its EXIF orientation, as the Capture window shows it
    if !decoded { errors[n] = sprint("its image cannot be decoded: %", decode_error); continue; }
    prepared[n] = .{image, which, true};
    array_add(*inputs, .{image.pixels, image.width, image.height, image.width * 4});
    array_add(*input_of, n);
  }
  if !inputs.count  return errors;

  results := segment_images(engine, inputs,, temp);
  for result, r: results {
    n := input_of[r];
    item := items[n];
    if result.failed { errors[n] = copy_string(result.error); continue; }
    p := prepared[n];
    segs, mask := segment_mcap_chunks(engine, result, p.which, p.image.width, p.image.height);
    stored: bool;
    store_error: string;
    for 1..OCR_STORE_ATTEMPTS {
      stored, store_error = mcap_store(item.path, front = .[segs], back = .[mask]);
      if stored  break;
      sleep_milliseconds(250);
    }
    if !stored { errors[n] = sprint("cannot be stored: %", store_error); continue; }
    entry, load_error := catalog_load(item.path);
    if entry {
      lock(*viewer.mutex);
      array_add(*viewer.catalog_updates, .{item.index, entry});
      unlock(*viewer.mutex);
    } else {
      errors[n] = sprint("was segmented, but could not be read back: %", load_error);
    }
  }
  return errors;
}

// -segment: segments the folders set to segment once the catalog has loaded, and ends the app when
// done (0, or 1 if any image failed).
segment_ui_segment_and_exit :: (busy: bool) {
  finish :: (code: s32, format: string, args: ..Any) {
    if code  log_error(format, ..args);
    else     log(format, ..args);
    app.exit_code = code;
    app.should_exit = true;
    app.segment_and_exit = false;
  }
  if viewer.scan_error { finish(1, "[Segment scan] %", viewer.scan_error); return; }
  if !segment_ui.scan_requested {
    // After the catalog, and after the engines made at startup, which use the same thread.
    if busy || viewer.catalog_loaded < viewer.captures.count || !ocr_ui.engine_tried || !labels_ui.engine_tried  return;
    segment_ui.scan_requested = true;
    segment_ui.window_open = true;   // its progress bar
    if !segment_scan_start(quiet = true)  finish(0, "[Segment scan] -segment: nothing to segment (no image in a folder set to segment needs it).");
    return;
  }
  scan := *ocr_ui.scan;
  if scan.running {
    done, total, failed, rate, left, elapsed := ocr_scan_progress();
    step := max(total / 20, 1);
    if done / step > segment_ui.logged_step {
      segment_ui.logged_step = done / step;
      log("[Segment scan] % segmented, % failed, % elapsed, % left, % images/s.", ocr_scan_progress_text(done, total), failed, ocr_duration_text(elapsed), ocr_duration_text(left), FF(rate, 1));
    }
    return;
  }
  if busy  return;
  finish(ifx scan.failed || segment_ui.engine_error then cast(s32) 1 else 0, "[Segment scan] -segment is done: % of % segmented in %, % failed.%",
         scan.done - scan.failed, scan.total, ocr_duration_text(scan.elapsed_seconds), scan.failed, ifx segment_ui.engine_error then tprint(" %", segment_ui.engine_error) else "");
}

// ---------------------------------------------------------------- the Segmentation window

segment_ui_frame :: () { using ImGui;
  busy := ocr_ui.thread_busy;
  if app.segment_and_exit  segment_ui_segment_and_exit(busy);
  segment_erase_popup();
  if !segment_ui.window_open  return;
  SetNextWindowSize(.{640, 520}, .FirstUseEver);
  if Begin(SEGMENT_WINDOW, *segment_ui.window_open) {
    segment_window_contents();
  }
  End();
}

segment_window_contents :: () { using ImGui;
  busy := ocr_ui.thread_busy;
  scan := *ocr_ui.scan;

  // The model.
  model_file := tprint("%/%.onnx", segment_models_folder(), SEGMENT_MODEL);
  if !busy && segment_ui.engine {
    Text("% on %", SEGMENT_MODEL, ifx segment_ui.engine.provider == .CUDA then "the GPU" else "the CPU");
    if segment_ui.engine.gpu_explanation  TextColored(app_theme.palette.dim, "%", segment_ui.engine.gpu_explanation);
  } else if !busy && segment_ui.engine_error {
    TextColored(app_theme.palette.warning, "% could not start: %", SEGMENT_MODEL, segment_ui.engine_error);
  } else if !file_exists(model_file) {
    TextColored(app_theme.palette.warning, "The model is not there (%): make it with tools/export_mask2former.py (see modules/Segment/README.md) and rebuild.", model_file);
  } else {
    Text("%", SEGMENT_MODEL);
    TextColored(app_theme.palette.dim, "Loads when a scan starts (about 2.4 GB of GPU memory).");
  }
  TextColored(app_theme.palette.dim, "Finds what is where in each image: people, animals, things, and stuff such as sky, wall and grass. Only in the folders set to segment, and never in videos.");
  Spacing();

  // The library folders.
  segment_update_counts();
  if BeginTable("folders", 4, .Borders | .RowBg | .SizingStretchProp) {
    TableSetupColumn("Library folder", .WidthStretch);
    TableSetupColumn("Segment", .WidthFixed);
    TableSetupColumn("Segmented", .WidthFixed);
    TableSetupColumn("", .WidthFixed);
    TableHeadersRow();
    for segment_ui.counts {
      PushID(cast(s32) it_index);
      TableNextRow();
      TableNextColumn();
      TextUnformatted(ifx it.folder then it.folder else "(the library's top)");
      TableNextColumn();
      on := library_folder_segments(it.folder);
      if busy && scan.segmenting  BeginDisabled();
      if Checkbox("##segment", *on)  library_folder_set_segment(it.folder, on);
      if busy && scan.segmenting  EndDisabled();
      if IsItemHovered(.AllowWhenDisabled)  SetTooltip(ifx on then "Scans read this folder's images" else "Scans leave this folder alone");
      TableNextColumn();
      if it.catalogued  Text("% of %", it.segmented, it.images);
      else              TextColored(app_theme.palette.dim, "% of % ...", it.segmented, it.images);
      TableNextColumn();
      can_erase := it.segmented > 0 && !busy;
      if !can_erase  BeginDisabled();
      if SmallButton("Erase") {
        free(segment_ui.erase_folder,, context.default_allocator);
        segment_ui.erase_folder = copy_string(it.folder,, context.default_allocator);
        segment_ui.erase_asked = true;
      }
      if !can_erase  EndDisabled();
      if IsItemHovered(.AllowWhenDisabled)  SetTooltip("Takes this folder's segments out of its files (asks first)");
      PopID();
    }
    EndTable();
  }
  gone := library_folders_gone();
  for gone {
    TextColored(app_theme.palette.dim, "\"%\" is set to segment, but the library has no folder by that name.", it);
    SameLine();
    PushID(temp_c_string(it));
    if SmallButton("Forget it")  library_folder_forget(it);
    PopID();
  }
  Spacing();

  // The scan.
  if scan.running && scan.segmenting {
    Text("Segmenting:");
    ocr_scan_progress_ui();
    return;
  }
  any_on := false;
  for segment_ui.counts  if library_folder_segments(it.folder)  any_on = true;
  catalog_ready := viewer.catalog_loaded >= viewer.captures.count;
  can_scan := !busy && catalog_ready && any_on;
  if !can_scan  BeginDisabled();
  if Button("Segment the folders set to segment")  segment_scan_start();
  if !can_scan  EndDisabled();
  if IsItemHovered(.AllowWhenDisabled) {
    if !any_on              SetTooltip("Tick Segment for a folder first");
    else if !catalog_ready  SetTooltip("Once the catalog has loaded");
    else if busy            SetTooltip("After what the OCR thread is doing (a scan, or loading a model)");
    else                    SetTooltip("Reads every image in those folders not yet segmented, storing each as it goes: a scan can stop and start again");
  }
  if scan.segmenting && scan.total && scan.done {
    TextColored(app_theme.palette.dim, "Last scan: % of % images in %, % failed.", scan.done, scan.total, ocr_duration_text(scan.elapsed_seconds), scan.failed);
  }
}

// The window's per-folder counts, worked out again only when the catalog has changed.
segment_update_counts :: () {
  if segment_ui.counted_generation == viewer.catalog_generation && segment_ui.counts.count == viewer.folders.count  return;
  segment_ui.counted_generation = viewer.catalog_generation;
  array_reset_keeping_memory(*segment_ui.counts);
  for viewer.folders  array_add(*segment_ui.counts, .{folder = it, catalogued = true});
  for viewer.captures {
    if it.removed  continue;
    count: *Segment_Folder_Count;
    for * c: segment_ui.counts  if c.folder == it.folder { count = c; break; }
    if !count  continue;
    entry := it.catalog;
    if !entry { count.catalogued = false; continue; }
    if entry.media == .VIDEO  continue;
    count.images += 1;
    if entry.segment_classes  count.segmented += 1;
  }
}

// ---------------------------------------------------------------- the Capture window's overlay
//
// The segments of the capture shown, over its image: each segment in its class's colour (a shade of
// it for each instance, so two people side by side differ), its edges darker, with the class under
// the mouse in a tooltip. A viewer worker reads the segs and mask chunks (never the images); the
// main thread makes the texture from the mask, again only when the capture, the opacity or the
// classes hidden change.

// A capture's segments as read for the overlay. Heap, with everything in it.
Segment_Overlay :: struct {
  index: s32 = -1;
  error: string;               // why there is none
  width, height: s32;          // the mask's
  map: [] u8;                  // a segment id per pixel; 0: none
  segments: [] Mcap_Segment;   // largest first
  by_id: [256] s16;            // into segments, for each id; -1: none
}

// On a viewer worker: the capture's segments by SEGMENT_MODEL, or why not.
segment_overlay_read :: (path: string) -> *Segment_Overlay {
  overlay := New(Segment_Overlay,, context.default_allocator);
  for * overlay.by_id  it.* = -1;
  m, original, ok, error := mcap_read_without_original(path,, temp);
  if !ok { overlay.error = copy_string(error,, context.default_allocator); return overlay; }
  segs := mcap_find_chunk(m.front_chunks, mcap_tag(MCAP_SEGMENTS_TAG), SEGMENT_MODEL);
  mask := mcap_find_chunk(m.back_chunks, mcap_tag(MCAP_MASK_TAG), SEGMENT_MODEL);
  if !segs || !mask { overlay.error = copy_string("it has not been segmented",, context.default_allocator); return overlay; }
  segments, segments_ok, segments_error := mcap_segments_parse(cast(string) segs.data,, temp);
  map, map_ok, map_error := mcap_mask_parse(mask.data,, temp);
  if !segments_ok || !map_ok { overlay.error = sprint("its segments cannot be read: %", ifx segments_error then segments_error else map_error,, context.default_allocator); return overlay; }

  push_allocator(context.default_allocator);
  overlay.width, overlay.height = map.width, map.height;
  overlay.map = array_copy(map.map);
  overlay.segments = NewArray(segments.segments.count, Mcap_Segment);
  for segments.segments {
    overlay.segments[it_index] = it;
    overlay.segments[it_index].class = copy_string(it.class);
    overlay.by_id[it.id] = cast(s16) it_index;
  }
  return overlay;
}

segment_overlay_free :: (overlay: *Segment_Overlay) {
  if !overlay  return;
  push_allocator(context.default_allocator);
  free(overlay.error);
  array_free(overlay.map);
  for overlay.segments  free(it.class);
  array_free(overlay.segments);
  free(overlay);
}

// Main thread, each frame: takes what the workers read, asks for the capture shown, and makes the
// texture when it is out of date.
segment_overlay_frame :: () {
  o := *segment_ui.overlay_state;
  lock(*viewer.mutex);
  done := viewer.segments_done;
  arrived := viewer.segments_done_request != 0 && viewer.segments_done_request == viewer.segments_request;
  viewer.segments_done = null;
  viewer.segments_done_request = 0;
  unlock(*viewer.mutex);
  if arrived {
    segment_overlay_free(o.overlay);
    o.overlay = done;
    o.loading = false;
    o.dirty = true;
  } else {
    segment_overlay_free(done);
  }

  // Nothing shown: let it all go.
  index := viewer.full_index;
  entry: *Catalog_Entry;
  if viewer.full_open && index >= 0 && index < viewer.captures.count && !viewer.captures[index].removed  entry = viewer.captures[index].catalog;
  wanted := o.show && entry && entry.segment_classes.count > 0;
  if !wanted {
    gpu_texture_release(*o.texture);
    if !viewer.full_open { segment_overlay_free(o.overlay); o.overlay = null; o.entry = null; o.loading = false; }
    return;
  }

  // The capture shown, read again when its catalog entry is new (stepped to, or segmented again).
  if o.entry != entry {
    o.entry = entry;
    o.loading = true;
    lock(*viewer.mutex);
    viewer.segments_wanted = index;
    viewer.segments_request += 1;
    unlock(*viewer.mutex);
    signal(*viewer.wake);
  }
  if o.dirty && o.overlay && o.overlay.index == index {
    o.dirty = false;
    segment_overlay_build();
  }
}

// The device went (graphics/DX11.jai): the texture with it; made again from the mask held.
segment_overlay_device_lost :: () {
  o := *segment_ui.overlay_state;
  gpu_texture_release(*o.texture);
  o.dirty = true;
}

// The row above the image: the toggle, the opacity, and the classes in it, each to hide or show
// (right click to search for it).
segment_overlay_controls :: (index: s32) { using ImGui;
  if index < 0 || index >= viewer.captures.count  return;
  entry := viewer.captures[index].catalog;
  if !entry || !entry.segment_classes  return;
  o := *segment_ui.overlay_state;

  if Checkbox("Segments", *o.show)  o.dirty = true;
  if IsItemHovered()  SetTooltip("What segmentation found in this image, drawn over it (%)", SEGMENT_MODEL);
  if !o.show { SameLine(); TextColored(app_theme.palette.dim, "% found", entry.segment_count); return; }
  SameLine();
  SetNextItemWidth(140);
  percent := o.opacity * 100;
  if SliderFloat("##opacity", *percent, 10, 90, "opacity %.0f")  { o.opacity = percent / 100; o.dirty = true; }
  if o.loading { SameLine(); TextColored(app_theme.palette.dim, "Loading..."); }
  if o.overlay && o.overlay.index == index && o.overlay.error { SameLine(); TextColored(app_theme.palette.warning, "%", o.overlay.error); }

  // The legend: each class once, largest first, with its colour and how many of it.
  right_edge := GetCursorScreenPos().x + GetContentRegionAvail().x;
  first := true;
  for segment, n: entry.segments {
    seen := false;
    for 0..n - 1  if entry.segments[it].class == segment.class { seen = true; break; }
    if seen  continue;
    count := 0;
    for entry.segments  if it.class == segment.class  count += 1;
    label := ifx count > 1 then tprint("% x%##legend%", segment.class, count, n) else tprint("%##legend%", segment.class, n);
    width := GetFrameHeight() + GetStyle().ItemInnerSpacing.x + GetFrameHeight() + GetStyle().ItemInnerSpacing.x + CalcTextSize(temp_c_string(label), null, true).x;
    if !first {
      SameLine();
      if GetCursorScreenPos().x + width > right_edge  NewLine();
    }
    first = false;

    // Its colour: click to choose another (kept for every capture), right click to go back to its own.
    PushID(temp_c_string(segment.class));
    base := segment_class_base_colour(segment.class);
    rgb := float.[base[0] / 255.0, base[1] / 255.0, base[2] / 255.0];
    if ColorEdit3("##colour", *rgb, .NoInputs | .NoLabel | .NoAlpha | .NoOptions | .PickerHueWheel) {
      segment_colour_set(segment.class, .[cast(u8) (rgb[0] * 255 + 0.5), cast(u8) (rgb[1] * 255 + 0.5), cast(u8) (rgb[2] * 255 + 0.5)], save = false);
    }
    if IsItemDeactivatedAfterEdit()  segment_colours_save();   // once the picker closes, not every frame of a drag
    chosen := false;
    for o.colours  if it.class == segment.class  chosen = true;
    if BeginPopupContextItem("##reset") {
      if MenuItem("Back to the default colour", null, false, chosen)  segment_colour_set(segment.class, .[0, 0, 0], default = true);
      EndPopup();
    }
    PopID();
    SameLine(0, GetStyle().ItemInnerSpacing.x);
    shown := !array_find(o.hidden, segment.class);
    if Checkbox(temp_c_string(label), *shown) {
      if shown  { found, at_hidden := array_find(o.hidden, segment.class); if found { free(o.hidden[at_hidden],, context.default_allocator); array_unordered_remove_by_index(*o.hidden, at_hidden); } }
      else      array_add(*o.hidden, copy_string(segment.class,, context.default_allocator));
      o.dirty = true;
    }
    if IsItemClicked(.Right)  viewer_set_search(tprint("has:\"%\"", segment.class));
    if IsItemHovered()  SetTooltip("Shows or hides %; right click to search has:\"%\"", segment.class, segment.class);
  }
}

// Over the image drawn from image_min to image_max: the overlay's texture, and the segment under the mouse.
segment_overlay_draw :: (index: s32, image_min: ImGui.ImVec2, image_max: ImGui.ImVec2) { using ImGui;
  o := *segment_ui.overlay_state;
  overlay := o.overlay;
  if !o.show || !overlay || overlay.index != index || !o.texture.handle  return;
  draw := GetWindowDrawList();
  draw.AddImage(draw, gpu_texture_ref(o.texture), image_min, image_max);

  if !IsItemHovered()  return;   // the image
  mouse := GetMousePos();
  w, h := image_max.x - image_min.x, image_max.y - image_min.y;
  if w <= 0 || h <= 0  return;
  x := cast(s32) ((mouse.x - image_min.x) / w * overlay.width);
  y := cast(s32) ((mouse.y - image_min.y) / h * overlay.height);
  if x < 0 || y < 0 || x >= overlay.width || y >= overlay.height  return;
  id := overlay.map[cast(s64) y * overlay.width + x];
  at := overlay.by_id[id];
  if !id || at < 0  return;
  segment := overlay.segments[at];
  if array_find(o.hidden, segment.class)  return;
  SetTooltip("%   score %   % of the image", segment.class, FF(segment.score, 2), join(tprint("%", FF(segment.area * 100, 1)), "%",, temp));
}

// A class's colour (the same in every capture), a shade of it for its instance'th segment in a
// capture: the next ones darker, then round again.
segment_class_colour :: (class: string, instance: s64) -> [3] u8 {
  base := segment_class_base_colour(class);
  shade := 1 - 0.18 * cast(float) (instance % 4);
  return .[cast(u8) (base[0] * shade), cast(u8) (base[1] * shade), cast(u8) (base[2] * shade)];
}

// The colour chosen for class in the legend; else its own, from a hash of its name.
segment_class_base_colour :: (class: string) -> [3] u8 {
  for segment_ui.overlay_state.colours  if it.class == class  return it.rgb;
  return segment_class_default_colour(class);
}

// A hue from a hash of the name, at a fixed saturation and brightness.
segment_class_default_colour :: (class: string) -> [3] u8 {
  hash := fnv1a_hash(class.data, class.count);
  hue := cast(float) (hash % 360) / 60.0;
  saturation := 0.65;
  value := 0.95;
  sector := cast(s32) hue;
  f := hue - sector;
  p := value * (1 - saturation);
  q := value * (1 - saturation * f);
  t := value * (1 - saturation * (1 - f));
  rgb: [3] float;
  if sector == {
    case 0; rgb = .[value, t, p];
    case 1; rgb = .[q, value, p];
    case 2; rgb = .[p, value, t];
    case 3; rgb = .[p, q, value];
    case 4; rgb = .[t, p, value];
    case;   rgb = .[value, p, q];
  }
  return .[cast(u8) (rgb[0] * 255), cast(u8) (rgb[1] * 255), cast(u8) (rgb[2] * 255)];
}

// The overlay as RGBA pixels at the mask's size (temp): each pixel in its segment's colour at the
// opacity, its edges darker and stronger; hidden classes and no segment left clear.
segment_overlay_pixels :: (overlay: *Segment_Overlay, hidden: [] string, opacity: float) -> [] u8 {
  w, h := overlay.width, overlay.height;
  fill, edge: [256] [4] u8;   // by segment id
  alpha := cast(u8) clamp(opacity * 255, 0, 255);
  edge_alpha := cast(u8) clamp(opacity * 255 * 1.6 + 40, 0, 255);
  for segment, n: overlay.segments {
    if array_find(hidden, segment.class)  continue;
    instance := 0;
    for 0..n - 1  if overlay.segments[it].class == segment.class  instance += 1;
    c := segment_class_colour(segment.class, instance);
    fill[segment.id] = .[c[0], c[1], c[2], alpha];
    edge[segment.id] = .[c[0] / 3, c[1] / 3, c[2] / 3, edge_alpha];
  }
  pixels := NewArray(cast(s64) w * h * 4, u8,, temp);
  for y: 0..h - 1 {
    row := cast(s64) y * w;
    for x: 0..w - 1 {
      i := row + x;
      id := overlay.map[i];
      if !id  continue;
      on_edge := (x > 0 && overlay.map[i - 1] != id) || (x + 1 < w && overlay.map[i + 1] != id) ||
                 (y > 0 && overlay.map[i - w] != id) || (y + 1 < h && overlay.map[i + w] != id);
      colour := ifx on_edge then edge[id] else fill[id];
      memcpy(pixels.data + i * 4, colour.data, 4);
    }
  }
  return pixels;
}

#scope_file

// The texture, from the mask held.
segment_overlay_build :: () {
  o := *segment_ui.overlay_state;
  gpu_texture_release(*o.texture);
  overlay := o.overlay;
  if !overlay || !overlay.map.count  return;
  pixels := segment_overlay_pixels(overlay, o.hidden, o.opacity);
  o.texture = gpu_texture_create(pixels.data, overlay.width, overlay.height);
}

#scope_export

// The confirmation for erasing one folder's segments, as a modal popup; at the top level, each frame.
segment_erase_popup :: () { using ImGui;
  POPUP :: "Erase segments?";
  if segment_ui.erase_asked && !IsPopupOpen(POPUP)  OpenPopup(POPUP);
  if !BeginPopupModal(POPUP, null, .AlwaysAutoResize)  return;
  folder := segment_ui.erase_folder;
  Text("Erase the segments of every image in %?", ifx folder then folder else "the library's top");
  TextColored(app_theme.palette.dim, "Each file holding them is rewritten without them; its images, text, OCR and labels stay as they are.");
  TextColored(app_theme.palette.warning, "Only segmenting the folder again brings them back.");
  Spacing();
  if Button("Erase") {
    segment_ui.erase_asked = false;
    CloseCurrentPopup();
    segment_ui.window_open = true;   // the progress bar is in the OCR window, opened by the erase
    ocr_ui.window_open = true;
    ocr_erase_start(.SEGMENTS, folder = folder, one_folder = true);
  }
  SameLine();
  if Button("Cancel") || IsKeyPressed(.Escape) {
    segment_ui.erase_asked = false;
    CloseCurrentPopup();
  }
  EndPopup();
}