// 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. 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 call at a time (a mutex), from any thread.

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);
}

// Runs the model on input, a float32 tensor of shape, and copies its output (float32 or float16) into
// output as float32. ok is false, with why, when it fails or gives a different number of values.
coreml_run :: (model: *Coreml_Model, input: [] float, shape: [] s64, output: [] float) -> ok: bool, error: string {
  lock(*model.mutex);
  defer unlock(*model.mutex);
  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.
  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);
  }
  strides := NewArray(shape.count, s64,, temp);
  step: s64 = 1;
  for < shape { strides[it_index] = step; step *= it; }
  if step != input.count  return 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 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 false, tprint("Core ML could not take the input: %", error_text(error));
  autorelease(features);

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

  out_value := (cast(func_value) objc_msgSend)(result, selector("featureValueForName:"), NSString.getTempString(model.output_name));
  out_array := ifx out_value then objc_msgSend_typed(out_value, selector("multiArrayValue")) else null;
  if !out_array  return false, tprint("the model gives no %", model.output_name);
  count := cast(s64) objc_msgSend_typed(out_array, selector("count"));
  if count != output.count  return false, tprint("the model gives % values, not %", count, output.count);
  func_type :: #type (*void, Selector) -> s64 #c_call;
  type := (cast(func_type) objc_msgSend)(out_array, selector("dataType"));
  // The output is packed (one row of the pooled vector here); dataPointer is read before the pool drains.
  data := objc_msgSend_typed(out_array, selector("dataPointer"));
  if type == MULTIARRAY_FLOAT32 {
    memcpy(output.data, data, count * size_of(float));
  } else if type == MULTIARRAY_FLOAT16 {
    half := cast(*u16) data;
    for 0..count - 1  output[it] = float16_to_float32(half[it]);
  } else {
    return false, tprint("the model's output is of type 0x%, not float32 or float16", formatInt(type, base = 16));
  }
  return true, "";
}

#scope_file

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

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";