// Full-vocabulary sampling for Top K = 0. Ordinary settings sample entirely on GPU.
export function sampleVocabulary(logits, settings, recent, position) {
  const values = Float32Array.from(logits),
    seen = new Set(recent.slice(-64));
  for (const id of seen)
    if (id < values.length)
      values[id] =
        values[id] < 0
          ? values[id] * settings.repeatPenalty
          : values[id] / settings.repeatPenalty;
  const order = Uint32Array.from({ length: values.length }, (_, i) => i).sort(
    (a, b) => values[b] - values[a] || a - b,
  );
  if (settings.temperature <= 0) return order[0];
  const weights = new Float64Array(order.length),
    max = values[order[0]];
  let sum = 0;
  for (let i = 0; i < order.length; i++) {
    weights[i] = Math.exp((values[order[i]] - max) / settings.temperature);
    sum += weights[i];
  }
  let mass = 0,
    count = 0;
  do {
    mass += weights[count++];
  } while (count < order.length && mass < settings.topP * sum);
  let hash =
    (Math.imul(
      (settings.seed ^ Math.imul(position, 2654435761)) >>> 0,
      747796405,
    ) +
      2891336453) >>>
    0;
  hash = Math.imul((hash >>> ((hash >>> 28) + 4)) ^ hash, 277803737) >>> 0;
  hash = ((hash >>> 22) ^ hash) >>> 0;
  const threshold = ((hash + 0.5) / 4294967296) * mass;
  let cumulative = 0;
  for (let i = 0; i < count; i++) {
    cumulative += weights[i];
    if (cumulative >= threshold) return order[i];
  }
  return order[count - 1];
}

// Exact speculative rejection sampling (Leviathan et al., 2023).
export function candidateDistribution(words, settings) {
  const floats = new Float32Array(words.buffer, words.byteOffset, words.length),
    k = Math.min(settings.topK || 1, words.length / 2);
  const values = [];
  let total = 0;
  for (let i = 0; i < k; i++) {
    const mass = Math.exp((floats[i * 2] - floats[0]) / settings.temperature);
    values.push([words[i * 2 + 1], mass]);
    total += mass;
  }
  let retained = 0,
    count = 0;
  do {
    retained += values[count++][1];
  } while (count < k && retained < settings.topP * total);
  return new Map(
    values.slice(0, count).map(([id, mass]) => [id, mass / retained]),
  );
}
export function drawDistribution(distribution, random) {
  let sum = 0,
    last;
  for (const [id, p] of distribution) {
    sum += p;
    last = id;
    if (random < sum) return id;
  }
  return last;
}
export function verifySampledDraft(
  target,
  draft,
  token,
  acceptRandom,
  resampleRandom,
) {
  const q = draft.get(token) || 0,
    p = target.get(token) || 0;
  if (q > 0 && acceptRandom < Math.min(1, p / q))
    return { accepted: true, token };
  const residual = new Map();
  let total = 0;
  for (const [id, p] of target) {
    const mass = Math.max(0, p - (draft.get(id) || 0));
    residual.set(id, mass);
    total += mass;
  }
  return {
    accepted: false,
    token: drawDistribution(
      total > 1e-12
        ? new Map([...residual].map(([id, p]) => [id, p / total]))
        : target,
      resampleRandom,
    ),
  };
}
export function randomStream(seed, position, stream) {
  let x =
    (seed ^ Math.imul(position, 2654435761) ^ Math.imul(stream, 2246822519)) >>>
    0;
  x = (Math.imul(x, 747796405) + 2891336453) >>> 0;
  x = Math.imul((x >>> ((x >>> 28) + 4)) ^ x, 277803737) >>> 0;
  return ((((x >>> 22) ^ x) >>> 0) + 0.5) / 4294967296;
}