// Core_ML: a Core ML model (an .mlpackage made by coremltools) run on one float tensor, on macOS.
// The models that run faster here than through onnxruntime on the CPU (tools/export_siglip_coreml.py
// says how much) use it; the rest stay with modules/onnxruntime.
//
//   model, ok, error := coreml_load(path, "pixel_values", "pooler_output");
//   ok, error = coreml_run(model, input, s64.[1, 3, 384, 384], output);   // output: the floats it gives
//   coreml_free(model);
//
// Loading compiles the package (Core ML's compileModelAtURL, a fraction of a second) into a temporary
// .mlmodelc, removed by coreml_free. What Core ML makes of a model for the GPU, the first time each
// size runs, is kept by macOS across runs (the very first run of the OCR models took 5 s more, later
// ones none). The model runs on the CPU and GPU, never the Neural Engine
// unless asked: with float16 weights it can be badly wrong there (SigLIP 2: cosine 0.72 to the CPU).
// A model runs one prediction at a time (a mutex), from any thread; outputs are read outside it.

Coreml_Units :: enum s64 {   // MLComputeUnits
  CPU_ONLY          :: 0;
  CPU_AND_GPU       :: 1;
  ALL               :: 2;
  CPU_AND_ENGINE    :: 3;
}

Coreml_Model :: struct {
  model: *NSObject;          // MLModel, retained
  compiled: *NSObject;       // the .mlmodelc's NSURL, retained; removed on free
  input_name, output_name: string;
  mutex: Mutex;
}

coreml_load :: (path: string, input_name: string, output_name: string, units := Coreml_Units.CPU_AND_GPU) -> *Coreml_Model, ok: bool, error: string {
  init_objective_c();   // before anything Objective-C, the pool included; once per process does it all
  pool := objc_new(NSAutoreleasePool);
  defer NSAutoreleasePool.drain(pool);

  if !file_exists(path)  return null, false, tprint("% is not there.", path);
  package := NSURL.fileURLWithPath(NSString.getTempString(path));

  error: *NSError;
  func_compile :: #type (*void, Selector, *NSURL, **NSError) -> *NSURL #c_call;
  compiled := (cast(func_compile) objc_msgSend)(objc_getClass("MLModel"), selector("compileModelAtURL:error:"), package, *error);
  if !compiled  return null, false, tprint("Core ML could not compile %: %", path, error_text(error));

  configuration := objc_init(objc_alloc(objc_getClass("MLModelConfiguration")));
  defer release(configuration);
  func_units :: #type (*void, Selector, s64) -> void #c_call;
  (cast(func_units) objc_msgSend)(configuration, selector("setComputeUnits:"), cast(s64) units);

  func_load :: #type (*void, Selector, *NSURL, *NSObject, **NSError) -> *NSObject #c_call;
  model := (cast(func_load) objc_msgSend)(objc_getClass("MLModel"), selector("modelWithContentsOfURL:configuration:error:"), compiled, configuration, *error);
  if !model {
    remove_compiled(compiled);
    return null, false, tprint("Core ML could not load %: %", path, error_text(error));
  }

  result := New(Coreml_Model,, context.default_allocator);
  result.model    = xx retain(model);
  result.compiled = xx retain(compiled);
  result.input_name  = copy_string(input_name,, context.default_allocator);
  result.output_name = copy_string(output_name,, context.default_allocator);
  init(*result.mutex);
  return result, true, "";
}

coreml_free :: (model: *Coreml_Model) {
  if !model  return;
  pool := objc_new(NSAutoreleasePool);
  defer NSAutoreleasePool.drain(pool);
  release(model.model);
  remove_compiled(cast(*NSURL) model.compiled);
  release(model.compiled);
  destroy(*model.mutex);
  free(model.input_name,, context.default_allocator);
  free(model.output_name,, context.default_allocator);
  free(model,, context.default_allocator);
}

// A run's output, as float32 values one after another (C order): a view of Core ML's own array when it
// is float32 and packed, else a packed copy. Valid until coreml_release_output.
Coreml_Output :: struct {
  data: *float;
  dims: [8] s64;
  dim_count: s64;
  count: s64;
  array: *NSObject;   // the MLMultiArray, retained
  copy: *float;       // when data is a copy: the default allocator's
}

// Runs the model on input, a float32 tensor of shape. ok is false, with why, when it fails.
coreml_predict :: (model: *Coreml_Model, input: [] float, shape: [] s64) -> Coreml_Output, ok: bool, error: string {
  result: Coreml_Output;
  init_objective_c();
  pool := objc_new(NSAutoreleasePool);
  defer NSAutoreleasePool.drain(pool);

  // The input, over the caller's floats: packed, so each stride is the product of the sizes after it.
  strides := NewArray(shape.count, s64,, temp);
  step: s64 = 1;
  for < shape { strides[it_index] = step; step *= it; }
  if step != input.count  return result, false, tprint("the input has % values, not % for its shape %", input.count, step, shape);

  error: *NSError;
  func_array_init :: #type (*void, Selector, *void, *NSObject, s64, *NSObject, *void, **NSError) -> *NSObject #c_call;
  array := (cast(func_array_init) objc_msgSend)(objc_alloc(objc_getClass("MLMultiArray")), selector("initWithDataPointer:shape:dataType:strides:deallocator:error:"),
                                                input.data, numbers(shape), MULTIARRAY_FLOAT32, numbers(strides), null, *error);
  if !array  return result, false, tprint("Core ML could not take the input: %", error_text(error));
  autorelease(array);

  func_value :: #type (*void, Selector, *NSObject) -> *NSObject #c_call;
  value := (cast(func_value) objc_msgSend)(objc_getClass("MLFeatureValue"), selector("featureValueWithMultiArray:"), array);
  key := NSString.getTempString(model.input_name);
  func_dictionary :: #type (*void, Selector, *NSObject, *NSObject) -> *NSObject #c_call;
  dictionary := (cast(func_dictionary) objc_msgSend)(objc_getClass("NSDictionary"), selector("dictionaryWithObject:forKey:"), value, key);
  func_provider :: #type (*void, Selector, *NSObject, **NSError) -> *NSObject #c_call;
  features := (cast(func_provider) objc_msgSend)(objc_alloc(objc_getClass("MLDictionaryFeatureProvider")), selector("initWithDictionary:error:"), dictionary, *error);
  if !features  return result, false, tprint("Core ML could not take the input: %", error_text(error));
  autorelease(features);

  func_predict :: #type (*void, Selector, *NSObject, **NSError) -> *NSObject #c_call;
  lock(*model.mutex);
  prediction := (cast(func_predict) objc_msgSend)(model.model, selector("predictionFromFeatures:error:"), features, *error);
  unlock(*model.mutex);
  if !prediction  return result, false, tprint("Core ML could not run the model: %", error_text(error));

  out_value := (cast(func_value) objc_msgSend)(prediction, selector("featureValueForName:"), NSString.getTempString(model.output_name));
  out_array := ifx out_value then cast(*NSObject) objc_msgSend_typed(out_value, selector("multiArrayValue")) else null;
  if !out_array  return result, false, tprint("the model gives no %", model.output_name);

  out_shape   := integers(cast(*NSObject) objc_msgSend_typed(out_array, selector("shape")));
  out_strides := integers(cast(*NSObject) objc_msgSend_typed(out_array, selector("strides")));
  if out_shape.count > result.dims.count || out_shape.count != out_strides.count  return result, false, tprint("the model's output has an unexpected shape %", out_shape);
  result.dim_count = out_shape.count;
  result.count = 1;
  for out_shape { result.dims[it_index] = it; result.count *= it; }

  func_type :: #type (*void, Selector) -> s64 #c_call;
  type := (cast(func_type) objc_msgSend)(out_array, selector("dataType"));
  if type != MULTIARRAY_FLOAT32 && type != MULTIARRAY_FLOAT16  return result, false, tprint("the model's output is of type 0x%, not float32 or float16", formatInt(type, base = 16));
  packed := true;
  expected: s64 = 1;
  for < out_shape { if out_strides[it_index] != expected  packed = false; expected *= it; }

  data := objc_msgSend_typed(out_array, selector("dataPointer"));
  result.array = xx retain(out_array);   // outlives the pool, and with it the data
  if type == MULTIARRAY_FLOAT32 && packed {
    result.data = cast(*float) data;
    return result, true, "";
  }
  // A copy, packed: float16 widened, and rows Core ML padded (GPU outputs may be) closed up.
  result.copy = cast(*float) alloc(result.count * size_of(float),, context.default_allocator);
  result.data = result.copy;
  index: [8] s64;
  for i: 0..result.count - 1 {
    offset: s64 = 0;
    for d: 0..result.dim_count - 1  offset += index[d] * out_strides[d];
    result.copy[i] = ifx type == MULTIARRAY_FLOAT32 then (cast(*float) data)[offset] else float16_to_float32((cast(*u16) data)[offset]);
    d := result.dim_count - 1;
    while d >= 0 { index[d] += 1; if index[d] < out_shape[d] break; index[d] = 0; d -= 1; }
  }
  return result, true, "";
}

coreml_release_output :: (out: *Coreml_Output) {
  if out.array  release(out.array);
  if out.copy   free(out.copy,, context.default_allocator);
  out.* = .{};
}

// coreml_predict, its output copied into output (which must hold exactly as many values).
coreml_run :: (model: *Coreml_Model, input: [] float, shape: [] s64, output: [] float) -> ok: bool, error: string {
  out, ok, error := coreml_predict(model, input, shape);
  if !ok  return false, error;
  defer coreml_release_output(*out);
  if out.count != output.count  return false, tprint("the model gives % values, not %", out.count, output.count);
  memcpy(output.data, out.data, out.count * size_of(float));
  return true, "";
}

#scope_file

MULTIARRAY_FLOAT32 :: 0x10000 | 32;   // MLMultiArrayDataTypeFloat32
MULTIARRAY_FLOAT16 :: 0x10000 | 16;   // MLMultiArrayDataTypeFloat16

numbers :: (values: [] s64) -> *NSObject {
  objects := NewArray(values.count, *NSObject,, temp);
  func_number :: #type (*void, Selector, s64) -> *NSObject #c_call;
  for values  objects[it_index] = (cast(func_number) objc_msgSend)(objc_getClass("NSNumber"), selector("numberWithLongLong:"), it);
  func_array :: #type (*void, Selector, **NSObject, u64) -> *NSObject #c_call;
  return (cast(func_array) objc_msgSend)(objc_getClass("NSArray"), selector("arrayWithObjects:count:"), objects.data, xx objects.count);
}

// An NSArray of NSNumbers, as integers. Temp.
integers :: (array: *NSObject) -> [] s64 {
  if !array  return .[];
  count := cast(s64) objc_msgSend_typed(array, selector("count"));
  result := NewArray(count, s64,, temp);
  func_at :: #type (*void, Selector, u64) -> *NSObject #c_call;
  func_value :: #type (*void, Selector) -> s64 #c_call;
  for 0..count - 1 {
    number := (cast(func_at) objc_msgSend)(array, selector("objectAtIndex:"), xx it);
    result[it] = (cast(func_value) objc_msgSend)(number, selector("longLongValue"));
  }
  return result;
}

error_text :: (error: *NSError) -> string {
  if !error  return "no reason given";
  return copy_string(to_string(NSError.localizedDescription(error)),, temp);
}

remove_compiled :: (url: *NSURL) {
  func :: #type (*void, Selector, *NSURL, **NSError) -> BOOL #c_call;
  (cast(func) objc_msgSend)(NSFileManager.defaultManager(), selector("removeItemAtURL:error:"), url, null);
}

float16_to_float32 :: (h: u16) -> float {
  sign := (cast(u32) h & 0x8000) << 16;
  exponent := (cast(u32) h >> 10) & 0x1F;
  mantissa := cast(u32) h & 0x3FF;
  bits: u32;
  if exponent == 0 {
    if mantissa == 0  bits = sign;
    else {   // subnormal: normalised
      e: u32 = 113;
      while !(mantissa & 0x400) { mantissa <<= 1; e -= 1; }
      bits = sign | (e << 23) | ((mantissa & 0x3FF) << 13);
    }
  } else if exponent == 31  bits = sign | 0x7F80_0000 | (mantissa << 13);
  else                      bits = sign | ((exponent + 112) << 23) | (mantissa << 13);
  return (cast(*float) *bits).*;
}

CoreML_Framework :: #library,system,link_always "CoreML";

#import "Basic";
#import "File";
#import "File_Utilities";
#import "Thread";
#import "Objective_C";
#import "Objective_C/AppKit";