// Labels: images and text in one space, from an image-text encoder (SigLIP 2), so an image can be
// scored against any phrase: a list of labels, or words typed into a search.
//
//   engine, ok, error := labels_create(models_folder, "siglip2-base-patch16-384");
//   images := labels_embed_images(engine, pictures);    // one unit vector per picture
//   phrases := labels_embed_texts(engine, .["code editor", "schematic"]);
//   score := labels_score(engine, image_vector, phrase_vector);   // 0 to 1
//
// The models are listed in models_folder/models.config (fetch_models.jai downloads them) and run
// with onnxruntime (modules/onnxruntime): the vision model on an NVIDIA GPU through CUDA when it
// can, else on the CPU; the text model, small inputs and only a few at a time, on the CPU.

Labels_Provider :: enum { CPU; CUDA; }

// Tightly packed or not: stride is the bytes from one row to the next.
Labels_Image :: struct {
  pixels: *u8;   // RGBA
  width, height, stride: s32;
}

Labels_Model_Info :: struct {
  name: string;
  title: string;
  vision, text, tokenizer: string;   // file names in the models folder
  image_size:  s32   = 384;
  max_tokens:  s32   = 64;
  logit_scale: float = 0;
  logit_bias:  float = 0;
  image_batch: s32   = 8;
}

Labels_Engine :: struct {
  model: Labels_Model_Info;
  provider: Labels_Provider;          // the vision model's
  gpu_explanation: string;            // why the GPU is not in use; "" when it is
  env: *OrtEnv;
  memory_info: *OrtMemoryInfo;
  vision, text: Ort_Session;
  tokenizer: *Gemma_Tokenizer;
  dimensions: s32;                    // of every vector
}

LABELS_MODELS_CONFIG :: "models.config";

// The models models.config lists, in its order; from the context allocator.
labels_list_models :: (models_folder: string) -> [] Labels_Model_Info, error: string {
  models: [..] Labels_Model_Info;
  path := tprint("%/%", models_folder, LABELS_MODELS_CONFIG);
  text, ok := read_entire_file(path, log_errors = false,, temp);
  if !ok  return models, tprint("Cannot read %.", path);

  model: *Labels_Model_Info;
  version_seen := false;
  line_number := 0;
  for raw: split(text, "\n",, temp) {
    line_number += 1;
    line := trim(labels_strip_comment(raw));
    if !line.count  continue;
    if begins_with(line, "@version") { version_seen = true; continue; }
    if !version_seen  return models, tprint("% does not start with @version 1.", path);
    if begins_with(line, "[model ") && ends_with(line, "]") {
      model = array_add(*models);
      model.* = .{};
      model.name = copy_string(trim(slice(line, 7, line.count - 8)));
      continue;
    }
    colon := find_index_from_left(line, #char ":");
    if colon < 0 || !model { log_error("[Labels] %:%: expected `key: value` under a [model NAME]. Ignored.", path, line_number); continue; }
    key   := trim(slice(line, 0, colon));
    value := trim(advance(line, colon + 1));
    if value.count >= 2 && value[0] == #char "\"" && value[value.count - 1] == #char "\""  value = slice(value, 1, value.count - 2);
    number, number_ok := string_to_float(value);
    if key == {
      case "title";       model.title     = copy_string(value);
      case "vision";      model.vision    = copy_string(value);
      case "text";        model.text      = copy_string(value);
      case "tokenizer";   model.tokenizer = copy_string(value);
      case "image_size";  if number_ok  model.image_size  = cast(s32) number;
      case "max_tokens";  if number_ok  model.max_tokens  = cast(s32) number;
      case "logit_scale"; if number_ok  model.logit_scale = number;
      case "logit_bias";  if number_ok  model.logit_bias  = number;
      case "image_batch"; if number_ok  model.image_batch = cast(s32) number;
      case;  log_error("[Labels] %:%: unknown key %. Ignored.", path, line_number, key);
    }
  }
  return models, "";
}

// An engine for the model named in models_folder/models.config; its vision model on CUDA with
// prefer_gpu when every library that needs is present, else on the CPU.
labels_create :: (models_folder: string, model_name: string, prefer_gpu := true) -> *Labels_Engine, ok: bool, error: string {
  push_allocator(context.default_allocator);
  loaded, load_error := ort_load();
  if !loaded  return null, false, load_error;

  models, models_error := labels_list_models(models_folder);
  if models_error  return null, false, models_error;
  model: *Labels_Model_Info;
  for * models  if it.name == model_name { model = it; break; }
  if !model  return null, false, tprint("% has no [model %].", LABELS_MODELS_CONFIG, model_name);

  engine := New(Labels_Engine);
  engine.model = model.*;
  api := ort.api;
  ok, error := ort_check(api.CreateEnv(ort.log_level, "labels", *engine.env));
  if !ok { labels_destroy(engine); return null, false, error; }
  ort_check(api.CreateCpuMemoryInfo(.OrtArenaAllocator, .OrtMemTypeDefault, *engine.memory_info));

  engine.tokenizer, ok, error = tokenizer_load(tprint("%/%", models_folder, model.tokenizer));
  if !ok { labels_destroy(engine); return null, false, error; }

  HINT :: "run `jai fetch_models.jai` in modules/Labels and rebuild";
  cores := cast(s32) max(get_number_of_processors(.ALL_LOGICAL), 1);
  vision_path := tprint("%/%", models_folder, model.vision);
  use_gpu := false;
  if prefer_gpu {
    available, explanation := ort_cuda_check("Labelling");
    use_gpu = available;
    if !available  engine.gpu_explanation = copy_string(explanation);
  } else {
    engine.gpu_explanation = copy_string("Labelling runs on the CPU, as it was asked to.");
  }
  if use_gpu {
    ort_gpu_lock();   // one CUDA engine made or released at a time (modules/onnxruntime/session.jai)
    engine.vision, ok, error = ort_create_session(engine.env, vision_path, true, 2, output = "pooler_output", missing_hint = HINT);
    if ok  ok, error = labels_warm_up(engine);
    if !ok  ort_release_session(*engine.vision);
    ort_gpu_unlock();
    if ok {
      engine.provider = .CUDA;
    } else {
      engine.gpu_explanation = sprint("Labelling runs on the CPU: the CUDA libraries were found, but onnxruntime could not start on the GPU:\n%", error);
      use_gpu = false;
    }
  }
  if !use_gpu {
    engine.vision, ok, error = ort_create_session(engine.env, vision_path, false, cores, output = "pooler_output", missing_hint = HINT);
    if !ok { labels_destroy(engine); return null, false, error; }
    engine.provider = .CPU;
  }
  engine.text, ok, error = ort_create_session(engine.env, tprint("%/%", models_folder, model.text), false, max(cores / 2, 1), output = "pooler_output", missing_hint = HINT);
  if !ok { labels_destroy(engine); return null, false, error; }

  // The vectors' size, from a run of the text model on an empty text.
  probe := labels_embed_texts(engine, .[""],, temp);
  if !engine.dimensions { labels_destroy(engine); return null, false, "The text model gave no vector."; }
  return engine, true, "";
}

labels_destroy :: (engine: *Labels_Engine) {
  if !engine  return;
  if engine.provider == .CUDA  ort_gpu_lock();
  ort_release_session(*engine.vision);
  if engine.provider == .CUDA  ort_gpu_unlock();
  ort_release_session(*engine.text);
  tokenizer_free(engine.tokenizer,, context.default_allocator);
  if engine.memory_info  ort.api.ReleaseMemoryInfo(engine.memory_info);
  if engine.env          ort.api.ReleaseEnv(engine.env);
  free(engine.gpu_explanation,, context.default_allocator);
  free(engine,, context.default_allocator);
}

labels_gpu_status :: (engine: *Labels_Engine) -> in_use: bool, explanation: string {
  return engine.provider == .CUDA, engine.gpu_explanation;
}

// One unit vector per image, engine.dimensions long, one after another; from allocator. An image
// that fails gives a zero vector, and error says why (the last error).
labels_embed_images :: (engine: *Labels_Engine, images: [] Labels_Image, allocator := context.allocator) -> [] float, error: string {
  size := engine.model.image_size;
  dims := engine.dimensions;
  result := NewArray(images.count * dims, float, initialized = true,, allocator);
  error: string;
  plane := cast(s64) size * size;
  batch := cast(s64) max(engine.model.image_batch, 1);
  input := NewArray(plane * 3 * batch, float,, context.default_allocator);
  defer array_free(input,, context.default_allocator);

  first := 0;
  while first < images.count {
    count := min(batch, images.count - first);
    for 0..count - 1  labels_prepare_image(images[first + it], size, input.data + plane * 3 * it);
    shape := s64.[count, 3, size, size];
    out, ok, run_error := ort_run_float(*engine.vision, engine.memory_info, array_view(input, 0, plane * 3 * count), shape);
    if !ok {
      error = tprint("%", run_error);
    } else {
      if out.dim_count == 2 && out.dims[0] == count && out.dims[1] == dims {
        for n: 0..count - 1  labels_normalize_into(out.data + n * dims, dims, result.data + (first + n) * dims);
      } else {
        error = tprint("The vision model gave an output of shape %, not [%, %].", array_view(out.dims, 0, out.dim_count), count, dims);
      }
      ort_release_output(*out);
    }
    first += count;
  }
  return result, error;
}

// One unit vector per text, as labels_embed_images gives them; from allocator.
labels_embed_texts :: (engine: *Labels_Engine, texts: [] string, allocator := context.allocator) -> [] float, error: string {
  if !texts.count  return .[], "";
  tokens := engine.model.max_tokens;
  ids := NewArray(texts.count * tokens, s64,, temp);
  for text, n: texts {
    encoded := tokenizer_encode(engine.tokenizer, text, tokens);
    memcpy(ids.data + n * tokens, encoded.data, tokens * size_of(s64));
  }
  shape := s64.[texts.count, tokens];
  out, ok, run_error := ort_run(*engine.text, engine.memory_info, ids.data, ids.count * size_of(s64), .ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64, shape);
  if !ok  return .[], tprint("%", run_error);
  defer ort_release_output(*out);
  if out.dim_count != 2 || out.dims[0] != texts.count  return .[], tprint("The text model gave an output of shape %.", array_view(out.dims, 0, out.dim_count));
  dims := cast(s32) out.dims[1];
  if !engine.dimensions  engine.dimensions = dims;
  result := NewArray(texts.count * dims, float,, allocator);
  for n: 0..texts.count - 1  labels_normalize_into(out.data + n * dims, dims, result.data + n * dims);
  return result, "";
}

// How well a phrase fits an image, 0 to 1, from their unit vectors: the model's own sigmoid.
labels_score :: (engine: *Labels_Engine, image: *float, text: *float) -> float {
  return labels_score(engine.model, image, text, engine.dimensions);
}

labels_score :: (model: Labels_Model_Info, image: *float, text: *float, dimensions: s64) -> float {
  cosine: float;
  for 0..dimensions - 1  cosine += image[it] * text[it];
  logit := exp(model.logit_scale) * cosine + model.logit_bias;
  return 1.0 / (1.0 + exp(-logit));
}

// A vector as half floats, for storing; and back. Temp / allocator.
labels_to_half :: (vector: [] float) -> [] u16 {
  result := NewArray(vector.count, u16,, temp);
  for vector  result[it_index] = labels_float_to_half(it);
  return result;
}

labels_from_half :: (half: [] u16, allocator := context.allocator) -> [] float {
  result := NewArray(half.count, float,, allocator);
  for half  result[it_index] = float16_to_float32(cast(Float16) it);
  return result;
}

// A float as 16-bit float bits: the Float16 module's float_to_half_fast3 (Fabian Giesen's, public
// domain), copied with #no_aoc. It relies on unsigned wraparound, which a build with arithmetic
// overflow checks (the app's release build) stops at; the module's own is not marked.
labels_float_to_half :: (value: float32) -> u16 #no_aoc {
  f32_infinity :: cast(u32) 255 << 23;
  f16_infinity :: cast(u32) 31  << 23;
  magic : float32 : 0h0780_0000;
  sign_mask  : u32 : 0x8000_0000;
  round_mask : u32 : 0xffff_f000;

  using bits: union { f: float32; u: u32; }
  f = value;
  sign := u & sign_mask;
  u ^= sign;
  half: u16;
  if u >= f32_infinity {
    half = ifx u > f32_infinity then cast(u16) 0x7e00 else 0x7c00;   // NaN to a quiet NaN, infinity to infinity
  } else {                                                           // normal, denormal or zero
    u &= round_mask;
    f *= magic;
    u -= round_mask;
    if u > f16_infinity  u = f16_infinity;                           // too large: infinity
    half = cast(u16) (u >> 13);
  }
  return half | cast(u16) (sign >> 16);
}

#scope_module

// Runs the vision model once, on a blank image, when it is made. As OCR's engines are (see
// ocr_warm_up): with cuBLASLt 13.8, the process's first CUDA session destroyed without ever having
// run left the next one to crash on its first run.
labels_warm_up :: (engine: *Labels_Engine) -> ok: bool, error: string {
  size := engine.model.image_size;
  input := NewArray(3 * size * size, float, initialized = true,, context.default_allocator);
  defer array_free(input,, context.default_allocator);
  out, ok, error := ort_run_float(*engine.vision, engine.memory_info, input, s64.[1, 3, size, size]);
  if !ok  return false, tprint("The first run on the GPU failed: %", error);
  ort_release_output(*out);
  return true, "";
}

// The image resized to size x size (not keeping its shape, as SigLIP's processor does) with PIL's
// bilinear filter, which averages over each output pixel's footprint when shrinking; then as planar
// RGB, scaled to -1..1, at dest.
labels_prepare_image :: (image: Labels_Image, size: s32, dest: *float) {
  // Horizontal pass into rows of size floats per channel, then vertical.
  across := labels_resample_weights(image.width, size);
  down   := labels_resample_weights(image.height, size);
  rows := NewArray(cast(s64) image.height * size * 3, float,, temp);
  for y: 0..image.height - 1 {
    row := image.pixels + cast(s64) y * image.stride;
    for x: 0..size - 1 {
      w := *across[x];
      r, g, b: float;
      for k: 0..w.weights.count - 1 {
        p := row + (w.first + k) * 4;
        r += p[0] * w.weights[k];  g += p[1] * w.weights[k];  b += p[2] * w.weights[k];
      }
      at := (cast(s64) y * size + x) * 3;
      rows[at] = r;  rows[at + 1] = g;  rows[at + 2] = b;
    }
  }
  plane := cast(s64) size * size;
  for y: 0..size - 1 {
    w := *down[y];
    for x: 0..size - 1 {
      r, g, b: float;
      for k: 0..w.weights.count - 1 {
        at := (cast(s64) (w.first + k) * size + x) * 3;
        r += rows[at] * w.weights[k];  g += rows[at + 1] * w.weights[k];  b += rows[at + 2] * w.weights[k];
      }
      i := cast(s64) y * size + x;
      dest[i]             = clamp(r, 0, 255) / 127.5 - 1;   // (v / 255 - 0.5) / 0.5
      dest[i + plane]     = clamp(g, 0, 255) / 127.5 - 1;
      dest[i + plane * 2] = clamp(b, 0, 255) / 127.5 - 1;
    }
  }
}

Resample_Weights :: struct { first: s32; weights: [] float; }

// For each of out positions, the input positions it takes and their weights: a triangle filter
// widened by the scale when shrinking, as PIL's bilinear resize. Temp.
labels_resample_weights :: (in: s32, out: s32) -> [] Resample_Weights {
  result := NewArray(out, Resample_Weights,, temp);
  scale := cast(float) in / out;
  filter_scale := max(scale, 1.0);
  support := filter_scale;   // the triangle reaches 1 unit, scaled
  for i: 0..out - 1 {
    centre := (i + 0.5) * scale;
    first := max(cast(s32) (centre - support + 0.5), 0);
    last  := min(cast(s32) (centre + support + 0.5), in);
    weights := NewArray(max(last - first, 1), float,, temp);
    total: float;
    for k: 0..last - first - 1 {
      x := (first + k - centre + 0.5) / filter_scale;
      w := max(0.0, 1 - abs(x));
      weights[k] = w;
      total += w;
    }
    if total > 0  for * weights  it.* /= total;
    else { first = clamp(cast(s32) centre, 0, in - 1); weights[0] = 1; }
    result[i] = .{first, weights};
  }
  return result;
}

labels_normalize_into :: (source: *float, dims: s64, dest: *float) {
  length: float;
  for 0..dims - 1  length += source[it] * source[it];
  length = sqrt(length);
  scale := ifx length > 0 then 1 / length else 0;
  for 0..dims - 1  dest[it] = source[it] * scale;
}

labels_strip_comment :: (line: string) -> string {
  in_quotes := false;
  for i: 0..line.count - 1 {
    c := line[i];
    if c == #char "\""  in_quotes = !in_quotes;
    else if c == #char "#" && !in_quotes && (i == 0 || line[i - 1] == #char " " || line[i - 1] == #char "\t")  return slice(line, 0, i);
  }
  return line;
}

#load "tokenizer.jai";

#import "Basic";
#import "String";
#import "File";
#import "Math";
#import "System";
#import "Hash_Table";
#import "Float16";
#import "onnxruntime";