// Labels (modules/Labels): the tokenizer's ids for known phrases (from Hugging Face's tokenizers on
// the same tokenizer.json), unit vectors from both encoders, and the model ranking a drawn text image
// as text rather than a photo. Needs the label models beside the program (the build copies them);
// without them, the suite says so and passes.

#scope_file

checks_run    := 0;
checks_failed := 0;

MAX_GROUPS :: 8;
group_names  : [MAX_GROUPS] string;
group_checks : [MAX_GROUPS] int;
group_fails  : [MAX_GROUPS] int;
group_count  := 0;

output_builder: *String_Builder;

emit :: (format_string: string, args: ..Any) {
  if output_builder  print_to_builder(output_builder, format_string, ..args);
  else               print(format_string, ..args);
} @PrintLike

group :: (name: string) {
  assert(group_count < MAX_GROUPS, "raise MAX_GROUPS");
  group_names[group_count] = name;
  group_count += 1;
  emit("\n-- % --\n", name);
}

expect :: (condition: bool, format: string, args: ..Any, loc := #caller_location) {
  checks_run += 1;
  if group_count  group_checks[group_count-1] += 1;
  if condition return;

  checks_failed += 1;
  if group_count  group_fails[group_count-1] += 1;
  emit("  FAIL  [%:%]  %\n", group_names[max(group_count-1, 0)], loc.line_number, tprint(format, ..args));
}

report :: () -> bool {
  emit("\n");
  for 0..group_count-1 {
    if group_fails[it]  emit("  % -- % checks, % FAILED\n", group_names[it], group_checks[it], group_fails[it]);
    else                emit("  % -- % checks\n", group_names[it], group_checks[it]);
  }
  if checks_failed  emit("\n% OF % LABELS CHECKS FAILED\n", checks_failed, checks_run);
  else              emit("\nALL LABELS TESTS PASSED  (% checks over % groups)\n", checks_run, group_count);
  return checks_failed == 0;
}

// From Hugging Face's tokenizers 0.23.2 on the same tokenizer.json, lowercased, cut and padded to 64.
Known_Tokens :: struct { text: string; ids: [] s64; }
KNOWN :: Known_Tokens.[
  .{"code editor",                    .[2564, 10476, 1]},
  .{"A screenshot of a terminal",     .[235250, 13703, 576, 476, 17145, 1]},
  .{"",                               .[1]},
];

// labels_float_to_half is a copy of the Float16 module's float_to_half_fast3 marked #no_aoc: the
// same bits for every value (here in a debug build, where the module's own is not stopped).
half_tests :: () {
  group("half floats");
  values: [..] float;
  values.allocator = temp;
  array_add(*values, ..float.[0, -0.0, 1, -1, 0.5, 65504, 65520, 1000000, -1000000, cast(float) 0.00006103515625, cast(float) 0.0000001, FLOAT32_INFINITY, -FLOAT32_INFINITY]);
  x: float = 0.00001;
  for 1..10_000 { array_add(*values, x, -x * 0.37); x *= 1.0027; }
  differ := 0;
  for values  if labels_float_to_half(it) != cast(u16) float32_to_float16(it)  differ += 1;
  expect(differ == 0, "% of % values convert to different bits", differ, values.count);
}

labels_tests :: () {
  half_tests();
  folder := labels_models_folder();
  if !file_exists(tprint("%/models.config", folder)) {
    emit("  (skipped: no label models in %)\n", folder);
    return;
  }
  group("engine");
  engine, created, error := labels_create(folder, LABELS_MODEL);
  expect(created, "the label engine starts: %", error);
  if !created  return;
  defer labels_destroy(engine);
  emit("  images on %; % dimensions\n", ifx engine.provider == .CUDA then "CUDA" else "the CPU", engine.dimensions);

  group("tokenizer");
  for KNOWN {
    ids := tokenizer_encode(engine.tokenizer, it.text, 64);
    same := ids.count == 64;
    for known, i: it.ids  if !same || ids[i] != known  same = false;
    for it.ids.count..63  if same && ids[it] != 0  same = false;   // <pad>
    expect(same, "\"%\" tokenizes as Hugging Face's does: % (got %)", it.text, it.ids, array_view(ids, 0, it.ids.count + 1));
  }

  group("vectors");
  prompts := string.["a screenshot of text on a white background", "a photo of a cat", "a photo of a mountain landscape"];
  texts, text_error := labels_embed_texts(engine, prompts,, temp);
  expect(!text_error && texts.count == prompts.count * engine.dimensions, "the prompts' vectors are made (%)", text_error);
  if text_error || texts.count != prompts.count * engine.dimensions  return;
  for 0..prompts.count - 1  expect(abs(vector_length(array_view(texts, it * engine.dimensions, engine.dimensions)) - 1) < 0.001, "the vector of \"%\" is a unit vector", prompts[it]);

  drawn, drawn_ok := test_drawn_text();
  if !drawn_ok {
    emit("  (the text image is drawn with GDI, on Windows only: skipped)\n");
    return;
  }
  image := Labels_Image.{drawn.pixels, drawn.width, drawn.height, drawn.stride};
  vector, image_error := labels_embed_images(engine, .[image],, temp);
  expect(!image_error && vector.count == engine.dimensions, "the image's vector is made (%)", image_error);
  if image_error || vector.count != engine.dimensions  return;
  expect(abs(vector_length(vector) - 1) < 0.001, "the image's vector is a unit vector");
  scores: [3] float;
  for 0..2  scores[it] = labels_score(engine, vector.data, texts.data + it * engine.dimensions);
  expect(scores[0] > scores[1] && scores[0] > scores[2], "drawn text scores as text (%) above a cat (%) and a landscape (%)", scores[0], scores[1], scores[2]);
}

vector_length :: (v: [] float) -> float {
  sum: float;
  for v  sum += it * it;
  return sqrt(sum);
}

#scope_export

run_labels_tests :: (builder: *String_Builder = null) -> bool {
  output_builder = builder;
  emit("=== labels test suite ===\n");
  push_allocator(context.default_allocator);
  labels_tests();
  return report();
}