import { yieldIfNeeded as p } from "../utilities/yielder.js";
import m from "../utilities/tokenParse.js";
import T, { SPECIALS as S } from "./BaseTokeniser.js";
function g(o, e) {
  return `${o}-::-${e}`;
}
function y(o) {
  const e = /* @__PURE__ */ new Map();
  for (let s = 0; s < o.length; s++) {
    const t = o[s];
    for (let n = 0; n < t.length - 1; n++) {
      const r = g(t[n], t[n + 1]), a = e.get(r) || {
        a: t[n],
        b: t[n + 1],
        count: 0,
        instances: /* @__PURE__ */ new Set()
      };
      a.count += 1, a.instances.add(s), e.set(r, a);
    }
  }
  return { pairs: e, tokens: o };
}
function f(o, e, s, t, n) {
  const r = g(e, s);
  if (o.pairs.has(r)) {
    const a = o.pairs.get(r);
    a.count += n, n > 0 ? a.instances.add(t) : a.count <= 0 ? o.pairs.delete(r) : a.instances.delete(t);
  } else
    o.pairs.set(r, { a: e, b: s, count: n, instances: /* @__PURE__ */ new Set([t]) });
}
function I(o) {
  let e = null, s = 0;
  for (const t of o.pairs.values())
    t.count > s && (s = t.count, e = t);
  return e;
}
function x(o, e) {
  return o.map((s) => {
    const t = [];
    for (let n = 0; n < s.length; n++)
      n < s.length - 1 && s[n] === e[0] && s[n + 1] === e[1] ? (t.push(e[0] + e[1]), n++) : t.push(s[n]);
    return t;
  });
}
function A(o, e) {
  e.instances.forEach((s) => {
    const t = o.tokens[s], n = [];
    for (let r = 0; r < t.length; r++)
      if (r < t.length - 1 && t[r] === e.a && t[r + 1] === e.b) {
        const a = e.a + e.b;
        n.push(a), r > 0 && (f(o, t[r - 1], e.a, s, -1), f(o, t[r - 1], a, s, 1)), r++, r < t.length - 1 && (f(o, e.b, t[r + 1], s, -1), f(o, a, t[r + 1], s, 1));
      } else
        n.push(t[r]);
    o.tokens[s] = n;
  }), o.pairs.delete(g(e.a, e.b));
}
class P extends T {
  targetSize;
  vocab = /* @__PURE__ */ new Set();
  vocabIndex = /* @__PURE__ */ new Map();
  merges = [];
  pretokenMap = /* @__PURE__ */ new Map();
  constructor(e, s) {
    super(), Array.isArray(e) ? (e.forEach((t, n) => {
      this.vocab.add(t), this.vocabIndex.set(t, n);
    }), s && (this.merges = s), this.targetSize = e.length, S.forEach((t) => {
      const n = e.indexOf(t);
      n !== -1 && this.addSpecialToken(t, n);
    })) : (this.addSpecialTokens(), this.targetSize = e);
  }
  addToken(e, s) {
    if (this.vocab.has(e))
      return this.vocabIndex.get(e);
    {
      this.vocab.add(e);
      const t = s !== void 0 ? s : this.vocab.size - 1;
      return this.vocabIndex.set(e, t), t;
    }
  }
  destroy() {
    this.vocab.clear(), this.vocabIndex.clear(), this.merges = [], this.pretokenMap.clear();
  }
  get trained() {
    return this.vocab.size > S.length && this.vocab.size <= this.targetSize;
  }
  get vocabSize() {
    return this.vocab.size;
  }
  get eosToken() {
    return this.vocabIndex.get("<eos>") ?? 0;
  }
  get bosToken() {
    return this.vocabIndex.get("<bos>") ?? 0;
  }
  get unkToken() {
    return this.vocabIndex.get("") ?? 1;
  }
  async train(e = [], s, t) {
    this.datasetID = t;
    let n = performance.now();
    const r = new Array(e.length);
    for (let i = 0; i < e.length; i++) {
      const h = e[i], l = new Array(h.length);
      for (let d = 0; d < h.length; d++)
        l[d] = m(h[d].content);
      n = await p(n, s, this.vocab.size), r[i] = l;
    }
    const a = r.flat(2), z = new Set(a);
    this.vocab = /* @__PURE__ */ new Set(), this.pretokenMap.clear(), this.merges = [], this.addSpecialTokens();
    const b = Array.from(z), v = b.map((i) => Array.from(i).map((l) => (this.vocab.add(l), l))), k = y(v);
    if (n = await p(n, s, this.vocab.size), this.vocab.size >= this.targetSize) {
      console.warn("Initial vocab size is greater than or equal to target size. No merges will be performed.");
      const i = /* @__PURE__ */ new Map();
      a.forEach((c) => {
        Array.from(c).forEach((u) => {
          i.set(u, (i.get(u) || 0) + 1);
        });
      });
      const h = Array.from(i.entries()).sort((c, u) => u[1] - c[1]);
      this.vocab = /* @__PURE__ */ new Set(), this.addSpecialTokens(), h.slice(0, this.targetSize - this.vocab.size).map(([c]) => c).forEach((c) => this.vocab.add(c)), this.vocabIndex.clear();
      let d = 0;
      for (const c of this.vocab.keys())
        this.vocabIndex.set(c, d++);
      return this.generateID(), this.emit("trainStatus", "trained"), this.vocab.size;
    }
    for (; this.vocab.size < this.targetSize && this.merges.length < this.targetSize; ) {
      const i = I(k);
      if (!i)
        break;
      this.merges.push([i.a, i.b]), this.vocab.add(i.a + i.b), A(k, i), n = await p(n, s, this.vocab.size);
    }
    b.forEach((i, h) => {
      const l = v[h];
      this.pretokenMap.set(i, l);
    }), this.vocabIndex.clear();
    let w = 0;
    for (const i of this.vocab.keys())
      this.vocabIndex.set(i, w++);
    return this.generateID(), this.emit("trainStatus", "trained"), this.vocab.size;
  }
  getVocab() {
    return Array.from(this.vocab);
  }
  getMerges() {
    return this.merges;
  }
  tokeniseWord(e) {
    let s = Array.from(e);
    return this.merges.forEach((t) => {
      s = x([s], t)[0];
    }), this.pretokenMap.set(e, s), s;
  }
  tokeniseStrings(e) {
    return e.map((s) => m(s).map((r) => this.pretokenMap.has(r) ? this.pretokenMap.get(r) : this.tokeniseWord(r)).flat(1));
  }
  tokenise(e, s) {
    const t = this.tokeniseStrings(e);
    return s ? t.map((n) => n.map((r) => this.vocabIndex.get(r) ?? this.unkToken)) : t.map((n) => n.map((r) => this.vocab.has(r) ? r : ""));
  }
  detokenise(e) {
    const s = this.getVocab();
    return e.map((n) => n.map((r) => s[r]).join(""));
  }
  encode(e) {
    return this.tokenise([e], !0)[0];
  }
  decode(e) {
    return this.detokenise([e])[0];
  }
}
export {
  P as default
};
