// AksharaLP v3 sentence rescoring (JS port of bhc/translit3.py search + bhc/agree.py). HARD guardrail: every output word is one
// of the candidates of its own slot; fixed slots pass through.
// lat: [[kind, [[cand, chLogSoftmax, wid], ...]], ...]  kind "x" (roman token) or "fix" (wid -2 = punct/digits/latin, else lexicon id)
const AUX = { "था": "ms", "थी": "fs", "थे": "mp", "थीं": "fp", "हूँ": "1s", "हूं": "1s", "है": "3s", "हैं": "pl", "हो": "2p" };
const PART = {};
for (const [w, g] of [["रहा", "ms"], ["रही", "fs"], ["रहे", "mp"], ["चुका", "ms"], ["चुकी", "fs"], ["चुके", "mp"], ["सका", "ms"], ["सकी", "fs"], ["सके", "mp"],
  ["गया", "ms"], ["गई", "fs"], ["गयी", "fs"], ["गए", "mp"], ["गये", "mp"], ["हुआ", "ms"], ["हुई", "fs"], ["हुए", "mp"],
  ["लिया", "ms"], ["ली", "fs"], ["लिए", "mp"], ["लिये", "mp"], ["दिया", "ms"], ["दी", "fs"], ["दिए", "mp"], ["दिये", "mp"],
  ["किया", "ms"], ["की", "fs"], ["किए", "mp"], ["किये", "mp"], ["आया", "ms"], ["आई", "fs"], ["आए", "mp"], ["आये", "mp"]]) PART[w] = g;
const SUBJ = { "मैं": "1s", "मै": "1s", "हम": "pl", "तुम": "2p", "आप": "hon", "तू": "2s" };
const BOUND = new Set(["कि", "और", "पर", "लेकिन", "तो", "जब", "क्योंकि", "अगर", "या", "मगर", "जो", ",", ".", "?", "!", "।"]);
const BAD_PA = new Set(["ms|fs", "ms|mp", "ms|fp", "ms|pl", "fs|ms", "fs|mp", "mp|ms", "mp|fs", "mp|1s", "mp|3s"]);
const SA = { "1s|1s": 1, "1s|3s": -1, "1s|pl": -1, "1s|2p": -1, "pl|pl": 1, "pl|1s": -1, "pl|3s": -1, "2p|2p": 1, "2p|1s": -1, "2p|3s": -1,
  "hon|pl": 1, "hon|1s": -1, "hon|3s": -1, "2s|3s": 1, "2s|1s": -1 };
const own = (o, k) => Object.prototype.hasOwnProperty.call(o, k);
function partClass(w) {
  if (own(PART, w)) return PART[w]; const n = [...w].length;
  if (n > 2 && w.endsWith("ता")) return "ms"; if (n > 2 && w.endsWith("ती")) return "fs"; if (n > 2 && w.endsWith("ते")) return "mp"; return null;
}
export function agreeFeats(words, i) {
  if (i === null || i <= 0) return [0, 0];
  const a = own(AUX, words[i]) ? AUX[words[i]] : null; if (a === null) return [0, 0];
  let pos = 0, neg = 0; const pc = partClass(words[i - 1]);
  if (pc !== null) { if (BAD_PA.has(pc + "|" + a)) neg++; else pos++; }
  for (let j = i - 1; j > Math.max(-1, i - 9); j--) {
    const u = words[j]; if (BOUND.has(u)) break;
    if (own(SUBJ, u)) { const v = SA[SUBJ[u] + "|" + a] || 0; pos += v > 0; neg += v < 0; break; }
  }
  return [pos, neg];
}
const NUM_RE = /^[0-9\u0966-\u096F]+$/;
export const V3DEF = { lat_k: 12, beam_s: 16, w_nlm: 0, w_agr: 0, w_agn: 0, nlm_margin: 99, nlm_k: 12, rec_n: 3 };
/** p: v2 params (w_lm, w_ch_s, oov_pen_s, use_lm) + V3DEF keys; lm: LexLM (logp(h, w), N, hasLM); nlm: NLM or null */
export function search3(lat, p, lm, nlm) {
  const useBi = p.use_lm && lm && lm.hasLM; if (!p.w_nlm) nlm = null;
  const agr = p.w_agr || p.w_agn;
  let beams = [[0, null, [], nlm ? nlm.initState() : null]];
  for (let [kind, opts] of lat) {
    if (kind === "fix") {
      const [t, , wid] = opts[0];
      if (wid === -2) {
        if (nlm && NUM_RE.test(t)) { const { ns } = nlm.scoreWords(beams.map((b) => b[3]), ["<n>"]); beams = beams.map(([s, h, ch], i) => [s, h, ch.concat([t]), ns[i][0]]); }
        else beams = beams.map(([s, h, ch, st]) => [s, h, ch.concat([t]), st]);
        continue;
      }
      opts = [[t, 0, wid]];
    }
    opts = opts.slice(0, p.lat_k); let cand = opts, lp = null, nst = null;
    if (nlm) { cand = opts.filter((o) => o[1] >= opts[0][1] - p.nlm_margin).slice(0, p.nlm_k); ({ lp, ns: nst } = nlm.scoreWords(beams.map((b) => b[3]), cand.map((c) => c[0]))); }
    const best = new Map();
    beams.forEach(([s, h, ch], i) => cand.forEach(([c, csc, wid], j) => {
      const w = wid === null || wid === undefined || wid === -1 ? null : wid;
      let tot = s + p.w_ch_s * csc;
      if (useBi) tot += p.w_lm * (lm.logp(h, w) - (w === null ? p.oov_pen_s : 0));
      if (nlm) tot += p.w_nlm * lp[i][j];
      const words = ch.concat([c]);
      if (agr) { const [ap, an] = agreeFeats(words, words.length - 1); tot += p.w_agr * ap - p.w_agn * an; }
      const key = !(nlm || agr) ? c : words.slice(-p.rec_n).join("\u0001");
      const cur = best.get(key);
      if (cur === undefined || tot > cur[0]) best.set(key, [tot, w, words, nlm ? nst[i][j] : null]);
    }));
    beams = [...best.values()].sort((a, b) => b[0] - a[0]).slice(0, p.beam_s);
  }
  const fin = (b) => b[0] + (useBi ? p.w_lm * lm.logp(b[1], lm.N) : 0) + (nlm ? p.w_nlm * nlm.eosLogp(b[3]) : 0);
  let bb = beams[0], bv = -Infinity; for (const b of beams) { const v = fin(b); if (v > bv) { bv = v; bb = b; } }
  return bb[2];
}