// The model part of the WhatDoWeSee Space, free of DOM code so it can be tested outside a browser.

/** Models the page can run. Each loads its files from its own Hub repo at a fixed revision. */
export const MODELS = [
  {
    id: "efficientnetv2_b0-v0.4.0",
    name: "EfficientNetV2-B0 · v0.4.0",
    repo: "StephanMeijer/WhatDoWeSee-EfficientNetV2-B0",
    revision: "v0.4.0",
  },
  {
    id: "mobilenetv4-v0.3.0",
    name: "MobileNetV4 · v0.3.0",
    repo: "StephanMeijer/WhatDoWeSee-MobileNetV4",
    revision: "v0.3.0",
  },
  {
    id: "efficientnet_b0-v0.3.0",
    name: "EfficientNet-B0 · v0.3.0",
    repo: "StephanMeijer/WhatDoWeSee-EfficientNet-B0",
    revision: "v0.3.0",
  },
  {
    id: "mobilenetv4-v0.2",
    name: "MobileNetV4 · v0.2",
    repo: "StephanMeijer/WhatDoWeSee-MobileNetV4",
    revision: "v0.2",
  },
  {
    id: "efficientnet_b0-v0.2",
    name: "EfficientNet-B0 · v0.2",
    repo: "StephanMeijer/WhatDoWeSee-EfficientNet-B0",
    revision: "v0.2",
  },
  {
    id: "efficientnet_b0-v0.1",
    name: "EfficientNet-B0 · v0.1",
    repo: "StephanMeijer/WhatDoWeSee-EfficientNet-B0",
    revision: "v0.1",
  },
];

export function fileUrl(model, file) {
  return `https://huggingface.co/${model.repo}/resolve/${model.revision}/${file}`;
}

/**
 * RGBA pixels already resized to size × size (no crop) to the model input: RGB scaled to [0, 1],
 * normalised with the config's mean and std, NCHW float32.
 */
export function toTensorData(rgba, size, mean, std) {
  const plane = size * size;
  const out = new Float32Array(3 * plane);
  for (let i = 0; i < plane; i++) {
    for (let c = 0; c < 3; c++) {
      out[c * plane + i] = (rgba[i * 4 + c] / 255 - mean[c]) / std[c];
    }
  }
  return out;
}

/** Every label with its probability and threshold, the present ones first, highest first. */
export function tags(logits, config) {
  return config.labels
    .map((id, i) => {
      const p = 1 / (1 + Math.exp(-logits[i]));
      const threshold = config.thresholds[id] ?? 0.5;
      return { id, name: config.names[i], definition: config.definitions[id], p, threshold, present: p >= threshold };
    })
    .sort((a, b) => Number(b.present) - Number(a.present) || b.p - a.p);
}

/** Run the model on prepared pixels; `ort` is onnxruntime-web (or -node in tests). */
export async function run(ort, session, config, rgba) {
  const size = config.input.shape[3];
  const data = toTensorData(rgba, size, config.input.mean, config.input.std);
  const output = await session.run({ [config.input.name]: new ort.Tensor("float32", data, [1, 3, size, size]) });
  return tags(output[config.output.name].data, config);
}