// cine-view-classifier — Phase 5 Session 1 community fixture.
//
// Classifies the active CINE slice into one of {4ch, lvot, sas, vla} via a
// 49 MB YOLOv11n classification ONNX. Hand-written single-file ESM (no JSX,
// no build step). Loaded by the install flow at runtime from a blob URL.
//
// Preprocessing follows DL-Modules/view_classification_yolov11/detector_onnx.py:
//   1. Per-image min-max normalize to [0, 1] float32
//   2. Replicate grayscale to 3 RGB channels
//   3. Circular pad to stride=32 multiple (NOT letterbox)
//   4. Transpose HWC -> CHW, add batch dim
//   5. Output is (1, 4) class probabilities; argmax over {4ch, lvot, sas, vla}.

const MODULE_ID = "cine-view-classifier";
const MODULE_VERSION = "0.1.0";
const WEIGHT_FILE = "no_other_view_yolo.onnx";
const WEIGHT_URL =
  "https://huggingface.co/JerryX/cmr-annotator-cine-view-classifier/resolve/main/checkpoints/no_other_view_yolo.onnx";
const WEIGHT_SHA256 =
  "6b6916e54b306655d56f892d1e1484acaf3b4e94a39cf410dccbaaa52cdb55fc";

const CLASS_NAMES = ["4ch", "lvot", "sas", "vla"];
const STRIDE = 32;

let sessionPromise = null;

function loadSession(host) {
  if (!sessionPromise) {
    sessionPromise = (async () => {
      host.meta.log.info("fetching view-classifier weights...");
      const buf = await host.weights.fetchCachedWeight(MODULE_ID, MODULE_VERSION, {
        file: WEIGHT_FILE,
        url: WEIGHT_URL,
        sha256: WEIGHT_SHA256,
      });
      host.meta.log.info("creating ONNX session...");
      return host.runtime.createOnnxSession(buf);
    })().catch((err) => {
      sessionPromise = null;
      throw err;
    });
  }
  return sessionPromise;
}

function normalizeSlice(slice) {
  const { pixels, width, height, bitsAllocated, pixelRepresentation } = slice;
  const n = width * height;
  const out = new Float32Array(n);
  let min = Infinity;
  let max = -Infinity;
  if (bitsAllocated === 16 && pixelRepresentation === 1) {
    const p = pixels;
    for (let i = 0; i < n; i++) {
      const v = p[i];
      if (v < min) min = v;
      if (v > max) max = v;
    }
    const range = max - min || 1;
    for (let i = 0; i < n; i++) out[i] = (p[i] - min) / range;
  } else {
    const p = pixels;
    for (let i = 0; i < n; i++) {
      const v = p[i];
      if (v < min) min = v;
      if (v > max) max = v;
    }
    const range = max - min || 1;
    for (let i = 0; i < n; i++) out[i] = (p[i] - min) / range;
  }
  return out;
}

function circularPadToStride(src, srcW, srcH, stride) {
  const newW = Math.ceil(srcW / stride) * stride;
  const newH = Math.ceil(srcH / stride) * stride;
  if (newW === srcW && newH === srcH) {
    return { data: src, width: srcW, height: srcH };
  }
  const out = new Float32Array(newW * newH);
  for (let y = 0; y < newH; y++) {
    const sy = y < srcH ? y : y % srcH;
    for (let x = 0; x < newW; x++) {
      const sx = x < srcW ? x : x % srcW;
      out[y * newW + x] = src[sy * srcW + sx];
    }
  }
  return { data: out, width: newW, height: newH };
}

function buildInputTensor(plane2d, w, h) {
  const planeSize = w * h;
  const tensor = new Float32Array(3 * planeSize);
  tensor.set(plane2d, 0);
  tensor.set(plane2d, planeSize);
  tensor.set(plane2d, 2 * planeSize);
  return tensor;
}

function softmax(logits) {
  let max = -Infinity;
  for (let i = 0; i < logits.length; i++) if (logits[i] > max) max = logits[i];
  const exps = new Float32Array(logits.length);
  let sum = 0;
  for (let i = 0; i < logits.length; i++) {
    exps[i] = Math.exp(logits[i] - max);
    sum += exps[i];
  }
  for (let i = 0; i < logits.length; i++) exps[i] /= sum || 1;
  return exps;
}

async function classifySlice(host, seriesUid, sliceIdx) {
  const slice = await host.data.getSlice(seriesUid, sliceIdx);
  const norm = normalizeSlice(slice);
  const padded = circularPadToStride(norm, slice.width, slice.height, STRIDE);
  const tensorData = buildInputTensor(padded.data, padded.width, padded.height);

  const session = await loadSession(host);
  const inputName = session.inputNames[0];
  const outputs = await session.run({
    [inputName]: {
      data: tensorData,
      dims: [1, 3, padded.height, padded.width],
    },
  });
  const out0 = outputs[session.outputNames[0]];
  let probs = out0.data;
  // Some YOLO11 classifier exports emit logits, others softmax-normalized
  // probabilities. If the values don't sum to ~1, run softmax ourselves.
  let sum = 0;
  for (let i = 0; i < probs.length; i++) sum += probs[i];
  if (Math.abs(sum - 1) > 0.01) probs = softmax(probs);

  let bestIdx = 0;
  for (let i = 1; i < probs.length; i++) {
    if (probs[i] > probs[bestIdx]) bestIdx = i;
  }
  return {
    label: CLASS_NAMES[bestIdx] ?? `class-${bestIdx}`,
    confidence: probs[bestIdx],
    probabilities: Array.from(probs).map((p, i) => ({
      label: CLASS_NAMES[i] ?? `class-${i}`,
      probability: p,
    })),
  };
}

export default function register(host) {
  // Track the current viewport's (seriesUid, sliceIdx) via slice-changed
  // events. The host doesn't expose a synchronous "give me the active series"
  // getter today (Phase-5 audit follow-up), so we keep our own latest-known.
  let active = null;
  const offSlice = host.events.on("slice-changed", (payload) => {
    active = payload;
  });

  const offRibbon = host.ui.registerRibbon({
    id: "classify-view",
    label: "Classify View",
    onClick: async (ctx) => {
      const series = active ?? null;
      if (!series) {
        ctx.host.ui.showToast(
          "Click a slice in a CINE viewport first, then try Classify View",
          "info",
        );
        return;
      }
      ctx.host.ui.showToast(
        `Classifying slice ${series.sliceIdx}…`,
        "info",
      );
      try {
        const result = await classifySlice(ctx.host, series.seriesUid, series.sliceIdx);
        const pct = Math.round(result.confidence * 100);
        ctx.host.ui.showToast(
          `View: ${result.label.toUpperCase()} (${pct}% confidence)`,
          "success",
        );
        ctx.host.meta.log.info(
          `classified slice ${series.sliceIdx}: ${result.label} (${pct}%)`,
          result.probabilities,
        );
      } catch (err) {
        const msg = err instanceof Error ? err.message : String(err);
        ctx.host.ui.showToast(`Classification failed: ${msg}`, "error");
        ctx.host.meta.log.error("classification failed", err);
      }
    },
  });

  host.meta.log.info("cine-view-classifier registered");

  return {
    host,
    dispose() {
      offRibbon();
      offSlice();
    },
  };
}