// segment-probe: what modules/Segment makes of image files. For each FILE, its segments (class, score,
// area); with -out FOLDER, the segments drawn over the image as a PNG, to look at; with -check FOLDER,
// how they compare with PyTorch's own segmentation of the same image (NAME.ref and NAME.ref.txt, from
// tools/export_mask2former.py --reference).
//
// Build: jai first.jai - segment_probe     Run: bin/segment-probe FILE... [-out FOLDER] [-check FOLDER] [-cpu] [-dump-input FOLDER]

main :: () {
  args := get_command_line_arguments();
  files: [..] string;
  out_folder, check_folder, dump_folder: string;
  models_folder := tprint("%segment-models", path_strip_filename(get_path_of_running_executable()));
  model_name := "mask2former-swin-base-coco-panoptic";
  prefer_gpu := true;
  i := 1;
  while i < args.count {
    if args[i] == "-out"    && i + 1 < args.count { out_folder    = args[i + 1]; i += 2; continue; }
    if args[i] == "-check"  && i + 1 < args.count { check_folder  = args[i + 1]; i += 2; continue; }
    if args[i] == "-models" && i + 1 < args.count { models_folder = args[i + 1]; i += 2; continue; }
    if args[i] == "-model"  && i + 1 < args.count { model_name    = args[i + 1]; i += 2; continue; }
    if args[i] == "-cpu" { prefer_gpu = false; i += 1; continue; }
    if args[i] == "-dump-input" && i + 1 < args.count { dump_folder = args[i + 1]; i += 2; continue; }
    array_add(*files, args[i]);
    i += 1;
  }
  if !files.count { print("Usage: segment-probe FILE... [-out FOLDER] [-check FOLDER] [-cpu] [-model NAME] [-models FOLDER]\n"); exit(2); }
  if out_folder  make_directory_if_it_does_not_exist(out_folder, recursive = true);

  start := current_time_monotonic();
  engine, ok, error := segment_create(models_folder, model_name, prefer_gpu);
  if !ok { print("Cannot start %: %\n", model_name, error); exit(1); }
  defer segment_destroy(engine);
  print("% on %, % classes (made in % ms)%\n", model_name, engine.provider, engine.classes.count, ms(start),
        ifx engine.gpu_explanation then tprint(": %", engine.gpu_explanation) else "");

  failures := 0;
  checked, matched, extra: s64;
  worst_iou: float64 = 1;
  agree, total_pixels, missed_confident: s64;
  model_ms, decode_ms: float64;
  done := 0;
  GROUP :: 8;   // images decoded and segmented together
  position := 0;
  while position < files.count {
    group := array_view(files, position, min(GROUP, files.count - position));
    position += group.count;

    images: [..] Image;
    names: [..] string;
    t := current_time_monotonic();
    for path: group {
      data, read_ok := read_entire_file(path, log_errors = false);
      if !read_ok { print("%: cannot be read\n", path); failures += 1; continue; }
      image, decoded, decode_error := image_decode(cast([] u8) data);   // turned by its EXIF orientation, as mcap shows it
      free(data);
      if !decoded { print("%: cannot decode: %\n", path, decode_error); failures += 1; continue; }
      array_add(*images, image);
      array_add(*names, path);
    }
    decode_ms += ms(t);
    if !images.count  continue;

    inputs := NewArray(images.count, Segment_Image,, temp);
    for images  inputs[it_index] = .{it.pixels, it.width, it.height, it.width * 4};
    if dump_folder {   // the model's input, as raw floats, to compare with the checkpoint's processor
      make_directory_if_it_does_not_exist(dump_folder, recursive = true);
      size := engine.model.image_size;
      floats := NewArray(3 * size * size, float,, temp);
      for inputs {
        segment_prepare(engine, it, floats.data);
        write_entire_file(tprint("%/%.input", dump_folder, path_filename(names[it_index])), floats.data, floats.count * size_of(float));
      }
    }
    t = current_time_monotonic();
    results := segment_images(engine, inputs,, temp);
    model_ms += ms(t);

    for result, n: results {
      path := names[n];
      image := images[n];
      done += 1;
      if result.failed { print("%: failed: %\n", path, result.error); failures += 1; continue; }
      print("% (% x %, map % x %): % segments\n", path_filename(path), image.width, image.height, result.width, result.height, result.segments.count);
      map_area := cast(float64) result.width * result.height;
      for result.segments  print("  %  %  %\n", formatFloat(it.score, trailing_width = 3), pad(engine.classes[it.class].name, 22), percent(it.pixels / map_area));

      if out_folder  write_overlay(engine, image, result, tprint("%/%.segments.png", out_folder, path_strip_extension(path_filename(path))));
      if check_folder {
        c := compare(engine, result, tprint("%/%.ref", check_folder, path_filename(path)), tprint("%/%.ref.txt", check_folder, path_filename(path)));
        if !c.ok { print("  check: %\n", c.error); failures += 1; continue; }
        print("  check: pixels agreeing on their class %; segments: % of % matched, worst IoU %%\n", percent(c.agree / cast(float64) c.pixels),
              c.matched, c.checked, formatFloat(c.worst, trailing_width = 4), ifx c.extra then tprint(", % extra", c.extra) else "");
        checked += c.checked; matched += c.matched; extra += c.extra; worst_iou = min(worst_iou, c.worst);
        agree += c.agree; total_pixels += c.pixels; missed_confident += c.missed_confident;
      }
    }
    for * images  image_free(it);
    reset_temporary_storage();
  }

  print("\n% images: decoding % ms each, segmenting % ms each (model and post-processing, batch %)\n", done,
        formatFloat(decode_ms / max(done, 1), trailing_width = 0), formatFloat(model_ms / max(done, 1), trailing_width = 0),
        ifx engine.provider == .CUDA then engine.model.batch_gpu else engine.model.batch_cpu);
  per := max(engine.stats.images, 1) * 1000.0;
  print("  of which: preparing % ms, the model % ms, making segments % ms, per image\n", formatFloat(engine.stats.prepare_us / per, trailing_width = 1),
        formatFloat(engine.stats.model_us / per, trailing_width = 1), formatFloat(engine.stats.post_us / per, trailing_width = 1));
  if check_folder && total_pixels {
    print("check: pixels agreeing on their class %; segments: % of % matched, % extra, worst IoU %\n",
          percent(agree / cast(float64) total_pixels), matched, checked, extra, formatFloat(worst_iou, trailing_width = 4));
    // Pass: nearly every pixel the same class, and every segment PyTorch is fairly sure of found. A
    // decoder or a rounding a step apart can tip a segment near a threshold either way.
    print("check: % confident segments (score at least %) missed\n", missed_confident, CONFIDENT);
    if cast(float64) agree < total_pixels * 0.995 || missed_confident  failures += 1;
  }
  exit(xx ifx failures then 1 else 0);
}

CONFIDENT :: 0.6;   // a segment PyTorch scores this or more must be matched

Comparison :: struct {
  ok: bool;
  error: string;
  pixels, agree: s64;            // pixels, and those given the same class (or none) by both
  checked, matched, extra: s64;  // PyTorch's segments, those with one of ours of the same class at IoU >= 0.5, and ours left over
  missed_confident: s64;         // PyTorch's segments scoring CONFIDENT or more that ours do not match
  worst: float64 = 1;            // the lowest IoU of a PyTorch segment with its best match
}

// Ours against PyTorch's: classes pixel by pixel, and each of its segments against our best of the same class.
compare :: (engine: *Segment_Engine, result: Segment_Result, ref_path: string, info_path: string) -> Comparison {
  c: Comparison;
  data, ok := read_entire_file(ref_path, log_errors = false,, temp);
  if !ok { c.error = tprint("no reference % (tools/export_mask2former.py --reference)", ref_path); return c; }
  if data.count < 8 || slice(data, 0, 4) != "M2FR" { c.error = tprint("% is not a reference", ref_path); return c; }
  w := (cast(s32) data[4]) + ((cast(s32) data[5]) << 8);
  h := (cast(s32) data[6]) + ((cast(s32) data[7]) << 8);
  if w != result.width || h != result.height { c.error = tprint("the reference is % x %, ours % x %", w, h, result.width, result.height); return c; }
  ref_map := cast([] u8) advance(data, 8);
  if ref_map.count != cast(s64) w * h { c.error = "the reference is cut short"; return c; }

  // Its segments' classes, by id.
  info, info_ok := read_entire_file(info_path, log_errors = false,, temp);
  if !info_ok { c.error = tprint("no %", info_path); return c; }
  ref_class: [256] s32;
  for * ref_class  it.* = -1;
  ref_score: [256] float64;
  ref_ids: [..] u8;
  ref_ids.allocator = temp;
  for line: split(info, "\n",, temp) {
    parts := split(trim(line), " ",, temp);
    if parts.count < 2  continue;
    id := string_to_int(parts[0]);
    ref_class[id] = cast(s32) string_to_int(parts[1]);
    if parts.count >= 3  ref_score[id] = string_to_float64(parts[2]);
    array_add(*ref_ids, cast(u8) id);
  }
  our_class: [256] s32;
  for * our_class  it.* = -1;
  for result.segments  our_class[it.id] = it.class;

  c.pixels = ref_map.count;
  for 0..ref_map.count - 1  if ref_class[ref_map[it]] == our_class[result.map[it]]  c.agree += 1;

  used: [256] bool;
  for rid: ref_ids {
    c.checked += 1;
    best: float64;
    best_id: u8;
    for result.segments {
      if it.class != ref_class[rid]  continue;
      both, either: s64;
      for p: 0..ref_map.count - 1 {
        a := ref_map[p] == rid;
        b := result.map[p] == it.id;
        if a && b  both += 1;
        if a || b  either += 1;
      }
      iou := ifx either then cast(float64) both / either else 1.0;
      if iou > best { best, best_id = iou, it.id; }
    }
    if best >= 0.5 { c.matched += 1; used[best_id] = true; }
    else if ref_score[rid] >= CONFIDENT  c.missed_confident += 1;
    if best < 0.5  print("  check: PyTorch's % (score %) has no match here (best IoU %)\n", engine.classes[ref_class[rid]].name, formatFloat(ref_score[rid], trailing_width = 3), formatFloat(best, trailing_width = 3));
    c.worst = min(c.worst, best);
  }
  for result.segments  if !used[it.id]  c.extra += 1;
  c.ok = true;
  return c;
}

// The segments in their classes' colours over the image, at the map's size, as a PNG.
write_overlay :: (engine: *Segment_Engine, image: Image, result: Segment_Result, path: string) {
  w, h := result.width, result.height;
  pixels := NewArray(cast(s64) w * h * 4, u8,, temp);
  for y: 0..h - 1  for x: 0..w - 1 {
    sx := min(cast(s32) ((x + 0.5) * image.width / w), image.width - 1);
    sy := min(cast(s32) ((y + 0.5) * image.height / h), image.height - 1);
    source := image.pixels + (cast(s64) sy * image.width + sx) * 4;
    dest := pixels.data + (cast(s64) y * w + x) * 4;
    id := result.map[cast(s64) y * w + x];
    if id {
      colour := segment_colour(result.segments[id - 1].class);
      for c: 0..2  dest[c] = cast(u8) ((cast(s32) source[c] + colour[c]) / 2);
    } else {
      for c: 0..2  dest[c] = source[c] / 3;   // no segment: darkened
    }
    dest[3] = 255;
  }
  stbi_write_png(temp_c_string(path), w, h, 4, pixels.data, w * 4);
}

ms :: (since: Apollo_Time) -> float64 {
  return to_float64_seconds(current_time_monotonic() - since) * 1000;
}

percent :: (fraction: float64) -> string {
  return join(tprint("%", formatFloat(fraction * 100, trailing_width = 1)), "%",, temp);
}

pad :: (text: string, width: s64) -> string {
  if text.count >= width  return text;
  builder: String_Builder;
  builder.allocator = temp;
  append(*builder, text);
  for 1..width - text.count  append(*builder, " ");
  return builder_to_string(*builder,, temp);
}

#import "Basic";
#import "String";
#import "File";
#import "File_Utilities";
#import "System";
#import "Image_Decode";
#import "Segment";
#import "stb_image_write";