// Proofbird GECToR grammar models in JavaScript - a line-by-line port of quickstart.py (the Python runner published with
// every model), so the browser demo returns exactly what the Python file returns. The model itself runs through whatever
// `run` function the caller passes in (onnxruntime-web on the page, onnxruntime-node in the parity test).
// Published by Patrick Andrew Hartnett - https://proofbird.app

const WORD = /[\p{L}\p{N}_]+(?:'[\p{L}\p{N}_]+)?|[^\p{L}\p{N}_\s]/gu; // a word (one apostrophe inside), or one mark
const ALNUM = /^[\p{L}\p{N}]$/u;
const CLOSING = ".,;:!?)]}%»”";

const isAlnum = (ch) => ch.length > 0 && ALNUM.test(Array.from(ch)[0]);

export class Gector {
  /**
   * files: { config, tokenizer, rules, verbs } - config.json and tokenizer.json parsed, language-rules.json parsed (or {}),
   * verb-form-vocab.txt as text (or ""). lang: "en" | "it" | "es" | "fr" | "de" | "pt-br" | "pt-pt". run(ids) -> { labels, d } (Float32Arrays,
   * [tokens x 5000] and [tokens x 2]).
   */
  constructor(files, lang, run, { keepConfidence = 0, minErrorProb = 0 } = {}) {
    this.lang = lang;
    this.run = run;
    this.keepConfidence = keepConfidence;
    this.minErrorProb = minErrorProb;
    const cfg = files.config;
    const n = Object.keys(cfg.id2label).length;
    this.labels = Array.from({ length: n }, (_, i) => cfg.id2label[String(i)]);
    this.keep = this.labels.indexOf(cfg.keep_label || "$KEEP");
    const dIds = Object.fromEntries(Object.entries(cfg.d_id2label).map(([k, v]) => [v, Number(k)]));
    this.incorrect = dIds[cfg.incorrect_label || "$INCORRECT"];
    const rules = files.rules || {};
    this.strip = rules.stripTrailingTerminator ? (rules.stripTerminators ?? ".!?") : "";
    const elision = rules.elision || {};
    this.glued = new Set(elision.nonElidingArticles || rules.nonElidingArticles || []);
    this.verbs = new Map();
    for (const line of (files.verbs || "").split("\n")) {
      if (!line.includes(":")) continue;
      const [words, tags] = line.trim().split(":");
      const [w1, w2] = words.split("_");
      this.verbs.set(`${w1}_${tags}`, w2);
    }
    this.opening = (lang === "es" ? "¿¡" : "") + "([{“«\"";
    this.tok = new WordPiece(files.tokenizer);
  }

  // --- one pass: words in, one tag per word out ---------------------------------------------------------------------
  async tag(words) {
    const { ids, wordIds } = this.tok.encode(words);       // [CLS] w1 ... wn [SEP]; [CLS] stands for $START
    const { labels, d } = await this.run(ids);
    const T = ids.length, L = this.labels.length;
    const result = [], seen = new Set();
    for (let t = 0; t < T; t++) {                          // one label per word, read at its first token
      const w = wordIds[t];
      const word = t === 0 ? 0 : (w === null ? null : w + 1);
      if (word === null || seen.has(word)) continue;
      seen.add(word);
      const probs = softmax(labels, t * L, L);
      probs[this.keep] += this.keepConfidence;
      let best = 0;
      for (let i = 1; i < L; i++) if (probs[i] > probs[best]) best = i;
      const wrong = softmax(d, t * 2, 2)[this.incorrect];
      const label = probs[best] >= this.minErrorProb ? this.labels[best] : "$KEEP";
      result.push([label, probs[best], wrong]);
    }
    if (Math.max(...result.map((r) => r[2])) < this.minErrorProb) {
      for (const r of result) r[0] = "$KEEP";             // no word is confidently wrong: keep the sentence
    }
    return { ids, result };
  }

  // --- apply one tag to one word ----------------------------------------------------------------------------------------
  edit(word, label) {
    if (label.startsWith("$APPEND_")) return word + " " + label.slice(8);
    if (word === "$START" || label === "$KEEP") return word;
    if (label.startsWith("$REPLACE_")) return label.slice(9);
    if (label === "$DELETE") return "";
    if (label.startsWith("$TRANSFORM_SUF_")) {             // change the ending: $TRANSFORM_SUF_<old>_<new>
      const rest = label.slice(15), cut = rest.indexOf("_");
      const old = rest.slice(0, cut), neu = rest.slice(cut + 1);
      if (!word.endsWith(old)) return word;
      const stem = word.slice(0, word.length - old.length);
      if (this.lang === "it" && neu.startsWith("h")) {     // spelling: an h only after a hard c/g
        const low = stem.toLowerCase();
        if (low.endsWith("chi") || low.endsWith("ghi")) return stem + neu.slice(1);  // vecchi+a -> vecchie
        if (!(low.endsWith("c") || low.endsWith("g"))) return word;
      }
      return stem + neu;
    }
    if (label.startsWith("$TRANSFORM_VERB_")) return this.verbs.get(`${word}_${label.slice(16)}`) ?? word;
    if (label === "$TRANSFORM_CASE_CAPITAL") return word.slice(0, 1).toUpperCase() + word.slice(1).toLowerCase();
    if (label === "$TRANSFORM_CASE_LOWER") return word.toLowerCase();
    if (label === "$TRANSFORM_CASE_UPPER") return word.toUpperCase();
    if (label === "$TRANSFORM_AGREEMENT_PLURAL") return word + "s";
    if (label === "$TRANSFORM_AGREEMENT_SINGULAR") return word.slice(0, -1);
    if (label === "$TRANSFORM_SPLIT_HYPHEN") return word.replaceAll("-", " ");
    if (label.startsWith("$MERGE_")) return word + " " + label;
    return word;
  }

  // --- the whole call: up to 5 passes, then the suggestions ---------------------------------------------------------
  async correct(text, maxPasses = 5) {
    let core = text, tail = "";
    const trimmed = text.trimEnd();
    if (this.strip && trimmed.length && this.strip.includes(trimmed.slice(-1))) {
      core = trimmed.slice(0, -1);
      tail = text.slice(trimmed.length - 1);
    }
    const plain = core.replaceAll("’", "'");               // the models were trained on straight apostrophes
    let words = this.words(plain);
    this.passes = [];
    for (let n = 1; n <= maxPasses; n++) {
      const { ids, result } = await this.tag(words);
      const src = ["$START", ...words];
      const m = Math.min(src.length, result.length);
      const edited = [];
      for (let i = 0; i < m; i++) edited.push(this.edit(src[i], result[i][0]));
      const tokens = edited.filter((e) => e).join(" ").replaceAll(" $MERGE_HYPHEN ", "-")
        .replaceAll(" $MERGE_SPACE ", "").split(/\s+/).filter((t) => t);
      const newWords = rejoin(tokens.filter((t) => t !== "$START")).flatMap((t) => this.words(t));
      const changed = !sameList(newWords, words);
      this.passes.push({
        pass: n, input: this.join(words.join(" ")), inputIds: ids,
        words: Array.from({ length: m }, (_, i) => ({
          word: src[i], tag: result[i][0], probability: result[i][1], wrong: result[i][2],
          edit: edited[i] !== src[i] ? edited[i] : null,
        })),
        result: this.join(newWords.join(" ")), changed,
      });
      if (!changed) break;
      words = newWords;
    }
    const found = this.suggestions(plain, words);          // same length as core, so positions match the real text
    found.push(...this.languageRules(core, found));        // German schon/schön, French ou/où (quickstart.py parity)
    found.sort((x, y) => x.start - y.start);
    let corrected = core;
    for (const f of [...found].reverse()) {                // apply right to left, so positions stay valid
      f.original = core.slice(f.start, f.start + f.length);
      corrected = corrected.slice(0, f.start) + f.replacement + corrected.slice(f.start + f.length);
    }
    return { corrected: corrected + tail, suggestions: found };
  }

  words(text) {
    return this.spans(text).map((s) => s[0]);
  }

  spans(text, surface = false) {
    const out = [];
    for (const m of text.matchAll(WORD)) {
      const g = m[0], start = m.index, end = start + g.length;
      const cut = g.indexOf("'");
      const article = cut < 0 ? g : g.slice(0, cut), rest = cut < 0 ? "" : g.slice(cut + 1);
      if (cut >= 0 && rest && this.glued.has(article.toLowerCase())) {
        const first = surface ? article + "'" : article;
        out.push([first, start, start + first.length], [rest, end - rest.length, end]);
      } else {
        out.push([g, start, end]);
      }
    }
    return out;
  }

  /** The rules the Proofbird app adds for shapes the model cannot learn without false alarms (model card, "What Proofbird
   * adds"), the same as quickstart.py: German "schon" -> "schön" at the end of a clause after a copula and a word like
   * trotzdem / sehr; French "ou" -> "où" after a negated knowing verb, an asking verb, "dis-moi" or "aucune idée" with a clause
   * after it, and as the last word of a statement with no comma before it. Skipped where the model already acted. */
  languageRules(text, found) {
    const free = (start, len) => !found.some((f) => f.start < start + len && start < f.start + Math.max(1, f.length));
    const out = [];
    if (this.lang === "de") {
      const re = /\b(?:ist|sind|war|waren|wäre|wären|bin|bist|seid|sei|seien|wird|werden|wurde|wurden|bleibt|bleiben|blieb|blieben|scheint|scheinen|schien|schienen|wirkt|wirken|wirkte|wirkten)\s+(?:(?:einfach|auch|doch|ja|eigentlich|halt)\s+)?(?:trotzdem|dennoch|wirklich|sehr|echt|total|ziemlich|besonders|unglaublich|wahnsinnig|richtig|wunderbar)\s+schon(?=\s*(?:[.!?,;:)»«"“”]|$))/gi;
      for (const m of text.matchAll(re)) {
        const start = m.index + m[0].length - 5;
        if (text.slice(start, start + 5) === "schon" && free(start, 5))
          out.push({ start, length: 5, original: "schon", replacement: "schön", rule: "schon/schön" });
      }
    } else if (this.lang === "fr") {
      const re = /\b(?:(?:ne|n['’])\s*(?:sais|sait|savons|savez|savent|savais|savait|savions|saviez|savaient|saurais|saurait)\s+(?:pas|plus|jamais|vraiment\s+pas|pas\s+trop|pas\s+du\s+tout)\s+|(?:me|te|se|nous|vous)\s+(?:demande|demandes|demandons|demandez|demandent|demandais|demandait|demandions|demandiez|demandaient|demander)\s+|(?:dis-moi|dites-moi|dis-nous|dites-nous|explique-moi|expliquez-moi|montre-moi|montrez-moi|devine|devinez)\s+|aucune\s+idée\s+)ou\s+(?=(?:je|j['’]|tu|il|elle|on|nous|vous|ils|elles|c['’]|ce|ça|cela|se|s['’]|l['’]|la|le|les|est|sont|était|étaient|sera|seront|va|vont|allait|aller|habite|habitent|habitait|vit|vivent|trouve|trouvent|trouver|mettre|ranger|acheter|vivre|dormir|manger|partir)(?![\wÀ-ÿ]))/gi;
      const starts = [];
      for (const m of text.matchAll(re)) {
        const start = m.index + m[0].length - m[0].match(/ou\s+$/i)[0].length;
        if (text.slice(start, start + 2) === "ou") starts.push(start);
      }
      const fin = text.match(/(\S)\s+(ou)\s*[.!…]*\s*$/i);
      if (fin && fin[2] === "ou" && !",;:".includes(fin[1]))
        starts.push(fin.index + fin[1].length + fin[0].slice(fin[1].length).search(/ou/i));
      for (const start of [...new Set(starts)].sort((x, y) => x - y))
        if (free(start, 2)) out.push({ start, length: 2, original: "ou", replacement: "où", rule: "ou/où" });
    }
    return out;
  }

  suggestions(original, correctedWords) {
    const spans = this.spans(original, true);
    const a = spans.map((s) => s[0]);
    const out = [];
    for (const [op, i1, i2, j1, j2] of new SequenceMatcher(a, correctedWords).getOpcodes()) {
      if (op === "equal") continue;
      let neu = this.join(correctedWords.slice(j1, j2).join(" "));
      let start, end;
      if (i1 === i2) {                                     // pure insertion
        if (i1 < a.length) {
          start = end = spans[i1][1];
          if (this.opening.includes(neu.slice(-1))) { /* an opening mark hugs the next word */ }
          else if (i1 === 0 || isAlnum(a[i1])) neu = neu + " ";
          else neu = " " + neu;
        } else {
          start = end = spans.length ? spans[spans.length - 1][2] : 0;
          neu = (isAlnum(neu.slice(0, 1)) ? " " : "") + neu;
        }
      } else {
        start = spans[i1][1]; end = spans[i2 - 1][2];
        if (neu && end < original.length && isAlnum(original[end]) && isAlnum(neu.slice(-1))) neu += " ";
        if (!neu) {                                        // pure deletion: take one neighbouring space with it
          if (end < original.length && original[end] === " ") end += 1;
          else if (start > 0 && original[start - 1] === " ") start -= 1;
        }
      }
      out.push({ start, length: end - start, original: original.slice(start, end), replacement: neu });
    }
    return out;
  }

  join(text) {
    // words back to normal text: no space before closing punctuation or after an opening mark
    const esc = (s) => s.replace(/[\\\]\[^-]/g, "\\$&");
    let t = text.split(/\s+/).filter((x) => x).join(" ");
    t = t.replace(new RegExp(` ([${esc(CLOSING)}])`, "gu"), "$1");
    return t.replace(new RegExp(`([${esc(this.opening.replace('"', ""))}]) `, "gu"), "$1");
  }
}

function rejoin(tokens) {                                  // a word the model ends with an apostrophe hugs the next word
  const out = [];
  for (const t of tokens) {
    const last = out[out.length - 1];
    if (last !== undefined && last.length > 1 && last.endsWith("'") && isAlnum(t.slice(0, 1))) out[out.length - 1] += t;
    else out.push(t);
  }
  return out;
}

function softmax(arr, offset, n) {
  let max = -Infinity;
  for (let i = 0; i < n; i++) if (arr[offset + i] > max) max = arr[offset + i];
  const out = new Float64Array(n);
  let sum = 0;
  for (let i = 0; i < n; i++) { out[i] = Math.exp(arr[offset + i] - max); sum += out[i]; }
  for (let i = 0; i < n; i++) out[i] /= sum;
  return out;
}

const sameList = (x, y) => x.length === y.length && x.every((v, i) => v === y[i]);

// --- BERT WordPiece over pre-split words (tokenizer.json: BertNormalizer + BertPreTokenizer + WordPiece) -------------
export class WordPiece {
  constructor(tj) {
    this.vocab = new Map(Object.entries(tj.model.vocab));
    this.unk = this.vocab.get(tj.model.unk_token);
    this.prefix = tj.model.continuing_subword_prefix ?? "##";
    this.maxChars = tj.model.max_input_chars_per_word ?? 100;
    const nz = tj.normalizer || {};
    this.clean = nz.clean_text !== false;
    this.chinese = nz.handle_chinese_chars !== false;
    this.lower = !!nz.lowercase;
    this.stripAccents = nz.strip_accents ?? this.lower;
    this.cls = this.vocab.get("[CLS]");
    this.sep = this.vocab.get("[SEP]");
  }

  encode(words) {
    const ids = [this.cls], wordIds = [null];
    words.forEach((w, wi) => {
      for (const piece of this.preTokenize(this.normalize(w))) {
        for (const id of this.wordPiece(piece)) { ids.push(id); wordIds.push(wi); }
      }
    });
    ids.push(this.sep); wordIds.push(null);
    return { ids, wordIds };
  }

  normalize(s) {
    let out = "";
    for (const ch of s) {
      const cp = ch.codePointAt(0);
      if (this.clean) {
        if (cp === 0 || cp === 0xfffd || (isControl(ch))) continue;
        if (/\s/u.test(ch)) { out += " "; continue; }
      }
      if (this.chinese && isCjk(cp)) { out += " " + ch + " "; continue; }
      out += ch;
    }
    if (this.stripAccents) out = out.normalize("NFD").replace(/\p{Mn}/gu, "");
    if (this.lower) out = out.toLowerCase();
    return out;
  }

  preTokenize(s) {
    const out = [];
    let cur = "";
    for (const ch of s) {
      if (/\s/u.test(ch)) { if (cur) out.push(cur); cur = ""; }
      else if (isPunct(ch)) { if (cur) out.push(cur); out.push(ch); cur = ""; }
      else cur += ch;
    }
    if (cur) out.push(cur);
    return out;
  }

  wordPiece(piece) {
    const chars = Array.from(piece);
    if (chars.length > this.maxChars) return [this.unk];
    const out = [];
    let start = 0;
    while (start < chars.length) {
      let end = chars.length, found = null;
      while (start < end) {
        let sub = chars.slice(start, end).join("");
        if (start > 0) sub = this.prefix + sub;
        if (this.vocab.has(sub)) { found = this.vocab.get(sub); break; }
        end -= 1;
      }
      if (found === null) return [this.unk];
      out.push(found);
      start = end;
    }
    return out;
  }
}

const isControl = (ch) => (ch === "\t" || ch === "\n" || ch === "\r") ? false : /\p{C}/u.test(ch);  // as HF tokenizers' is_other
const isPunct = (ch) => /^[!-\/:-@\[-`{-~]$/.test(ch) || /\p{P}/u.test(ch);
const isCjk = (cp) => (cp >= 0x4e00 && cp <= 0x9fff) || (cp >= 0x3400 && cp <= 0x4dbf) || (cp >= 0x20000 && cp <= 0x2a6df)
  || (cp >= 0x2a700 && cp <= 0x2b73f) || (cp >= 0x2b740 && cp <= 0x2b81f) || (cp >= 0x2b820 && cp <= 0x2ceaf)
  || (cp >= 0xf900 && cp <= 0xfaff) || (cp >= 0x2f800 && cp <= 0x2fa1f);

// --- Python's difflib.SequenceMatcher (no junk, autojunk off), enough for get_opcodes ---------------------------------
export class SequenceMatcher {
  constructor(a, b) {
    this.a = a; this.b = b;
    this.b2j = new Map();
    b.forEach((x, j) => { if (!this.b2j.has(x)) this.b2j.set(x, []); this.b2j.get(x).push(j); });
  }

  findLongestMatch(alo, ahi, blo, bhi) {
    let besti = alo, bestj = blo, bestsize = 0;
    let j2len = new Map();
    for (let i = alo; i < ahi; i++) {
      const next = new Map();
      for (const j of this.b2j.get(this.a[i]) || []) {
        if (j < blo) continue;
        if (j >= bhi) break;
        const k = (j2len.get(j - 1) || 0) + 1;
        next.set(j, k);
        if (k > bestsize) { besti = i - k + 1; bestj = j - k + 1; bestsize = k; }
      }
      j2len = next;
    }
    return [besti, bestj, bestsize];
  }

  getMatchingBlocks() {
    const la = this.a.length, lb = this.b.length;
    const queue = [[0, la, 0, lb]], blocks = [];
    while (queue.length) {
      const [alo, ahi, blo, bhi] = queue.pop();
      const [i, j, k] = this.findLongestMatch(alo, ahi, blo, bhi);
      if (k) {
        blocks.push([i, j, k]);
        if (alo < i && blo < j) queue.push([alo, i, blo, j]);
        if (i + k < ahi && j + k < bhi) queue.push([i + k, ahi, j + k, bhi]);
      }
    }
    blocks.sort((x, y) => x[0] - y[0] || x[1] - y[1] || x[2] - y[2]);
    const out = [];
    let [i1, j1, k1] = [0, 0, 0];
    for (const [i2, j2, k2] of blocks) {
      if (i1 + k1 === i2 && j1 + k1 === j2) k1 += k2;
      else { if (k1) out.push([i1, j1, k1]); [i1, j1, k1] = [i2, j2, k2]; }
    }
    if (k1) out.push([i1, j1, k1]);
    out.push([la, lb, 0]);
    return out;
  }

  getOpcodes() {
    let i = 0, j = 0;
    const out = [];
    for (const [ai, bj, size] of this.getMatchingBlocks()) {
      const tag = i < ai && j < bj ? "replace" : i < ai ? "delete" : j < bj ? "insert" : "";
      if (tag) out.push([tag, i, ai, j, bj]);
      i = ai + size; j = bj + size;
      if (size) out.push(["equal", ai, i, bj, j]);
    }
    return out;
  }
}