import { env } from "@huggingface/transformers";
import { readFileSync } from "node:fs";
import { OpenJev, choice, noul } from "/Users/shreyas/work/rnd/gliner2/work/open-jev/dist/index.js";

env.allowRemoteModels = false;
env.allowLocalModels = true;
env.localModelPath = "/Users/shreyas/work/rnd/gliner2/work/models";

const ref = JSON.parse(readFileSync("/Users/shreyas/work/rnd/gliner2/work/export/conversion/reference.json", "utf8"));
const dtype = process.argv[2] ?? "fp32";
const t0 = Date.now();
const jev = await OpenJev.load({ model: "gliner2-decide", device: "cpu", dtype });
console.log("loaded", jev.runtime, "in", Date.now() - t0, "ms");

// Re-implement enough of the family to check the token ids too.
let worst = 0, agree = 0, n = 0, idMismatch = 0;
const times = [];
for (const c of ref) {
  const questions = c.questions.map(([task, labels, descs]) =>
    labels.length === 2 && labels.includes("yes") && labels.includes("no")
      ? { kind: "noul-like", q: choice(task, labels, descs ?? undefined) }
      : { kind: "choice", q: choice(task, labels, descs ?? undefined) });
  const qs = questions.map((x) => x.q);
  const t1 = Date.now();
  const answers = await jev.decide(c.text, qs);
  times.push(Date.now() - t1);
  answers.forEach((a, i) => {
    const labels = c.questions[i][1];
    const pyProbs = c.probabilities[i];
    const jsProbs = labels.map((l) => a.probabilities[l]);
    const d = Math.max(...jsProbs.map((p, k) => Math.abs(p - pyProbs[k])));
    worst = Math.max(worst, d);
    const pyBest = labels[pyProbs.indexOf(Math.max(...pyProbs))];
    agree += a.choice === pyBest ? 1 : 0; n += 1;
    console.log(`  ${c.questions[i][0].slice(0, 30).padEnd(30)} js=${a.choice.padStart(34)} ${a.confidence.toFixed(4)}  py=${pyBest.padStart(34)} ${Math.max(...pyProbs).toFixed(4)}  |dp|=${d.toExponential(2)}`);
  });
}
console.log(`dtype=${dtype} argmax agreement ${agree}/${n}, worst |dp| ${worst.toExponential(2)}, median decide ${times.sort((a,b)=>a-b)[Math.floor(times.length/2)]} ms (node cpu)`);
await jev.dispose();