// Ports of the two model repos' prompt.py files: questions in, token ids out, probabilities back.

// Python's json.dumps(x, ensure_ascii=False) with its default ", " / ": " separators.
export function pyDumps(x) {
  if (x === null || x === undefined) return "null";
  if (typeof x === "string") return JSON.stringify(x);
  if (typeof x === "number" || typeof x === "boolean") return JSON.stringify(x);
  if (Array.isArray(x)) return "[" + x.map(pyDumps).join(", ") + "]";
  return "{" + Object.entries(x).map(([k, v]) => JSON.stringify(k) + ": " + pyDumps(v)).join(", ") + "}";
}

// ------------------------------------------------------------------ d1-omni-600M (prompt.py)

export const QTYPES = { choice: 0, score: 1, noul: 2 };
const DELIM = { state: "<|reserved_7|>", q: "<|reserved_8|>", opt: "<|reserved_9|>", opt_end: "<|reserved_10|>", decide: "<|reserved_11|>" };
const MARKER = "<|mask|>";

export const escapeSpecial = (t) => t.replace(/<\|([A-Za-z0-9_]+)\|>/g, "<¦$1¦>");
const criterion = (v) => (typeof v === "string" ? v : pyDumps(v));
export const nOptions = (q) => (q.type === "noul" ? 2 : q.type === "choice" ? Object.keys(q.criteria).length : q.criteria.length);

function renderOptions(q, nounDefault, audio) {
  if (q.type === "choice") {
    const entries = Object.entries(q.criteria);
    if (audio) return entries.map(([k, v], i) => `option_${String(i).padStart(3, "0")}: ${criterion(v == null || v === "" ? k : v)}`);
    return entries.map(([k, v]) => (v == null || v === "" ? k : `${k}: ${criterion(v)}`));
  }
  if (q.type === "score") return q.criteria.map((c, i) => `level ${i}: ${criterion(c)}`);
  if (audio) return ["false: no", "true: yes"];
  const crit = q.criteria || nounDefault || {};
  const f = crit.false ?? crit.no, t = crit.true ?? crit.yes;
  return ["false: " + (f != null && f !== "" ? criterion(f) : "no, the statement does not hold"),
          "true: " + (t != null && t !== "" ? criterion(t) : "yes, the statement holds")];
}

export function omniEncode(tok, state, q, maxLen, nounDefault = null, audio = false, perOption = 24) {
  const id = (s) => {
    const ids = tok.encode(s, { add_special_tokens: false });
    if (ids.length !== 1) throw new Error(`not a single token: ${s}`);
    return ids[0];
  };
  const enc = (s) => tok.encode(escapeSpecial(s), { add_special_tokens: false });
  const opts = renderOptions(q, nounDefault, audio);
  const budget = Math.max(96, Math.min(opts.length * perOption + 32, Math.floor(maxLen / 2)));
  const per = Math.max(2, Math.floor((budget - 3 * opts.length) / opts.length));
  let question = [id(DELIM.q), ...enc(q.instructions)].slice(0, Math.max(16, budget));
  let markers = [];
  for (const text of opts) {
    markers.push(question.length + 1);
    question = question.concat([id(DELIM.opt), id(MARKER)], enc(" " + text).slice(0, per), [id(DELIM.opt_end)]);
  }
  question.push(id(DELIM.decide));
  const room = Math.max(0, maxLen - question.length - 2);
  const stateIds = [id(DELIM.state), ...enc(typeof state === "string" ? state : pyDumps(state)).slice(0, room)];
  const ids = [tok.bos_token_id ?? 1, ...stateIds, ...question].slice(0, maxLen);
  markers = markers.map((m) => m + 1 + stateIds.length);
  if (markers[markers.length - 1] >= maxLen) throw new Error("the options do not fit in the context");
  return { ids, markers };
}

export function temperatureKey(q) {
  const k = nOptions(q);
  return `${q.type}:` + (k <= 2 ? "2" : k <= 5 ? "3-5" : k <= 10 ? "6-10" : "11+");
}

// ------------------------------------------------------------------ d1-3B (prompt.py)

const FALLBACK_POOL = (() => {
  const out = [];
  for (let c = 65; c <= 90; c++) out.push(String.fromCharCode(c));
  for (let i = 0; i < 100; i++) out.push(String(i).padStart(2, "0"));
  for (let c = 97; c <= 122; c++) out.push(String.fromCharCode(c));
  for (let i = 0; i < 200; i++) out.push(`#${i}`);
  for (let a = 65; a <= 90; a++) for (let b = 65; b <= 90; b++) out.push(String.fromCharCode(a) + String.fromCharCode(b));
  return out;
})();

function optionCodes(labels) {
  const labs = labels.map((x) => String(x).trim());
  if (labs.length && labs.every((k) => k.length === 1 && /\p{L}/u.test(k))) return labs;
  if (labs.length <= 26) return labs.map((_, i) => String.fromCharCode(65 + i));
  return labs.map((_, i) => String(i).padStart(2, "0"));
}

function aliases(tok, labels) {
  const used = new Set(), out = [];
  const take = (raw) => {
    const e = tok.encode(raw, { add_special_tokens: false });
    if (e.length !== 1 || used.has(e[0])) return false;
    out.push([raw, e[0]]);
    used.add(e[0]);
    return true;
  };
  for (const code of optionCodes(labels)) {
    if (take(code)) continue;
    if (!FALLBACK_POOL.some(take)) throw new Error(`no single-token alias left for ${labels.length} options`);
  }
  return out;
}

function singleIds(tok, texts) {
  const out = [], seen = new Set();
  for (const t of texts) {
    const e = tok.encode(t, { add_special_tokens: false });
    if (e.length === 1 && !seen.has(e[0])) { out.push(e[0]); seen.add(e[0]); }
  }
  return out;
}

export function readoutIds3b(tok, q) {
  if (q.type === "noul") return [singleIds(tok, ["yes", "Yes", "YES"]), singleIds(tok, ["no", "No", "NO"])];
  if (q.type === "score") return q.criteria.map((_, i) => singleIds(tok, [String(i)]));
  return aliases(tok, Object.keys(q.criteria)).map(([code, tid]) => [tid, ...singleIds(tok, [` ${code}`]).filter((i) => i !== tid)]);
}

function questionBlock3b(tok, q) {
  if (q.type === "choice") {
    const labels = Object.keys(q.criteria);
    const codes = aliases(tok, labels);
    const lines = labels.map((lab, i) => `${codes[i][0]} ${q.criteria[lab] || lab.replace(/_/g, " ")}`).join("\n");
    return `${q.instructions}\n\nOptions:\n${lines}\n\nReply with the option code only.`;
  }
  if (q.type === "noul") {
    const extra = q.criteria ? `\nYes: ${q.criteria.true}\nNo: ${q.criteria.false}` : "";
    return `${q.instructions}${extra}\n\nReply with yes or no only.`;
  }
  const legend = q.criteria.map((name, i) => `${i} ${name}`).join("\n");
  return `${q.instructions}\n\n${legend}\n\nReply with a single digit 0-${q.criteria.length - 1} only.`;
}

export function render3b(tok, state, q, imageMarkup = "") {
  const bos = "<|startoftext|>";
  let body = "";
  if (state !== null && state !== undefined) {
    const block = typeof state === "string" ? `${state}\n\n` : JSON.stringify(state, null, 2) + "\n\n";
    body = `${block}\nQUESTION:\n`;
  }
  return `${bos}<|im_start|>user\n${imageMarkup}${body}${questionBlock3b(tok, q)}<|im_end|>\n<|im_start|>assistant\n`;
}

// ------------------------------------------------------------------ answers (api.py / prompt.py `answer`)

export function softmax(z) {
  const m = Math.max(...z);
  const e = z.map((v) => Math.exp(v - m));
  const s = e.reduce((a, b) => a + b, 0);
  return e.map((v) => v / s);
}

export function answer(q, probs) {
  if (q.type === "noul") return { type: "noul", noul: probs[0] };
  let best = 0;
  probs.forEach((p, i) => { if (p > probs[best]) best = i; });
  if (q.type === "choice") {
    const names = Object.keys(q.criteria);
    return { type: "choice", choice: names[best], confidence: probs[best], probabilities: Object.fromEntries(names.map((n, i) => [n, probs[i]])) };
  }
  return {
    type: "score", score: probs.reduce((a, p, i) => a + i * p, 0), confidence: probs[best],
    probabilities: Object.fromEntries(probs.map((p, i) => [String(i), p])),
    legend: Object.fromEntries(q.criteria.map((c, i) => [String(i), criterion(c)])),
  };
}