export function applyRepPenalty(
  logits: Float32Array,
  generatedTokens: number[],
  penalty: number,
): Float32Array {
  const next = new Float32Array(logits)
  for (const token of new Set(generatedTokens)) {
    next[token] = next[token] > 0 ? next[token] / penalty : next[token] * penalty
  }
  return next
}

export function applyMinP(logits: Float32Array, minP: number): Float32Array {
  const next = new Float32Array(logits)
  let maxLogit = -Infinity
  for (const value of next) {
    if (value > maxLogit) {
      maxLogit = value
    }
  }

  const probs = new Float64Array(next.length)
  let total = 0
  let maxProb = 0
  for (let index = 0; index < next.length; index += 1) {
    const prob = Math.exp(next[index] - maxLogit)
    probs[index] = prob
    total += prob
    if (prob > maxProb) {
      maxProb = prob
    }
  }

  const threshold = (maxProb / total) * minP
  for (let index = 0; index < next.length; index += 1) {
    if (probs[index] / total < threshold) {
      next[index] = -1e9
    }
  }

  return next
}

export function sampleWithTemperature(
  logits: Float32Array,
  temperature: number,
): number {
  let maxLogit = -Infinity
  for (const value of logits) {
    if (value > maxLogit) {
      maxLogit = value
    }
  }

  const probs = new Float64Array(logits.length)
  let total = 0
  for (let index = 0; index < logits.length; index += 1) {
    const prob = Math.exp(logits[index] / temperature - maxLogit / temperature)
    probs[index] = prob
    total += prob
  }

  let threshold = Math.random() * total
  for (let index = 0; index < probs.length; index += 1) {
    threshold -= probs[index]
    if (threshold <= 0) {
      return index
    }
  }

  return probs.length - 1
}