/**
 * UI-side previews of what a watermark does to a single decoding step.
 * These are display helpers; the authoritative computation happens in the
 * worker's generation loop.
 */

/**
 * Probabilities after adding `delta` to green logits, renormalized.
 *
 * Exact renormalization needs the full vocabulary. We only ship the top-N
 * candidates to the UI, so the un-shown tail is approximated by its expected
 * behaviour: a fraction `gamma` of the tail mass is green and gets e^delta.
 * With the top-N holding most of the mass this is visually indistinguishable
 * from the exact curve, and the shape (green up, red down) is exact.
 */
export function biasedProbs(
  probs: number[],
  isGreen: boolean[],
  delta: number,
  gamma: number,
): number[] {
  const e = Math.exp(delta);
  let shown = 0;
  let weighted = 0;
  for (let i = 0; i < probs.length; i++) {
    shown += probs[i];
    weighted += probs[i] * (isGreen[i] ? e : 1);
  }
  const tail = Math.max(0, 1 - shown);
  const z = weighted + tail * (1 - gamma + gamma * e);
  if (z <= 0) return probs.slice();
  return probs.map((p, i) => (p * (isGreen[i] ? e : 1)) / z);
}

/** Gumbel-max score R^(1/p), returned as log for numeric stability. */
export function gumbelLogScore(r: number, p: number): number {
  if (p <= 0) return -Infinity;
  return Math.log(r) / p;
}