"""Batch-one NumPy beam state for the pinned TrOCR generation policy.

Implements Transformers4.57's finished/running beam rules. Equal-score top-k
ties use lower indices, explicitly; Torch CPU/CUDA tie order is not portable.
Source equivalence and complete OCR qualification must be measured separately.
"""
import numpy as np


def top_indices(scores, count):
    scores = np.asarray(scores)
    assert scores.ndim == 1 and 1 <= count <= scores.size
    assert not np.isnan(scores).any() and not np.isposinf(scores).any() and np.isfinite(scores).sum() >= count
    threshold = np.partition(scores, scores.size - count)[scores.size - count]
    better = np.flatnonzero(scores > threshold)
    tied = np.flatnonzero(scores == threshold)[:count - len(better)]
    indices = np.concatenate([better, tied])
    return indices[np.lexsort((indices, -scores[indices]))]


def log_softmax(scores):
    scores = np.asarray(scores, dtype=np.float32)
    shifted = scores - scores.max(axis=-1, keepdims=True)
    return shifted - np.log(np.exp(shifted).sum(axis=-1, keepdims=True))


class TrOCRBeamState:
    def __init__(self, beams=10, max_new_tokens=256, length_penalty=1., suppress_tokens=()):
        if beams < 2 or not 1 <= max_new_tokens <= 512 or length_penalty < 0:
            raise ValueError('Unsupported beam search settings')
        self.beams = beams;self.maximum = max_new_tokens + 1;self.penalty = length_penalty
        self.suppressed = tuple(suppress_tokens)
        if any(type(t) is not int or t < 0 for t in self.suppressed) or len(set(self.suppressed)) != len(self.suppressed):
            raise ValueError('Suppressed token IDs must be unique nonnegative integers')
        self.current = 1
        self.running = np.full((beams, self.maximum), 1, np.int64);self.running[:, 0] = 2
        self.running_scores = np.full(beams, -1e9, np.float32);self.running_scores[0] = 0
        self.finished = self.running.copy();self.finished_scores = np.full(beams, -1e9, np.float32)
        self.finished_flags = np.zeros(beams, bool);self.finished_lengths = np.zeros(beams, np.int64)
        self.done = False

    def advance(self, logits):
        logits = np.asarray(logits, dtype=np.float32)
        if self.done or logits.ndim != 2 or logits.shape[0] != self.beams or logits.shape[1] < 3 or not np.isfinite(logits).all():
            raise ValueError('Invalid scores or completed beam search')
        classes = logits.shape[1]
        if self.suppressed and max(self.suppressed) >= classes:raise ValueError('Suppressed ID outside vocabulary')
        logprobs = log_softmax(logits)
        # HF applies suppression after log-softmax, without renormalization.
        logprobs[:, self.suppressed] = -np.inf
        accumulated = logprobs + self.running_scores[:, None]
        choices = top_indices(accumulated.reshape(-1), 2 * self.beams)
        scores = accumulated.reshape(-1)[choices]
        parents, tokens = choices // classes, choices % classes
        candidates = self.running[parents].copy();candidates[:, self.current] = tokens
        stopped = (tokens == 2) | (self.current + 1 == self.maximum)

        live_scores = scores + stopped.astype(np.float32) * np.float32(-1e9)
        live = top_indices(live_scores, self.beams)
        self.running = candidates[live].copy();self.running_scores = live_scores[live].copy()

        eligible = stopped & (np.arange(2 * self.beams) < self.beams)
        normalized = scores / np.float32(self.current ** self.penalty)
        normalized = normalized + (~eligible).astype(np.float32) * np.float32(-1e9)
        merged_scores = np.concatenate([self.finished_scores, normalized])
        selected = top_indices(merged_scores, self.beams)
        self.finished = np.concatenate([self.finished, candidates])[selected].copy()
        self.finished_scores = merged_scores[selected].copy()
        self.finished_flags = np.concatenate([self.finished_flags, eligible])[selected].copy()
        self.finished_lengths = np.concatenate([self.finished_lengths, np.full(2 * self.beams, self.current, np.int64)])[selected].copy()

        self.current += 1
        possible = self.running_scores[0] / np.float32((self.current - 1) ** self.penalty)
        worst = self.finished_scores.min() if self.finished_flags.all() else np.float32(-1e9)
        self.done = bool(possible <= worst or stopped.all())
        return parents[live].copy(), tokens[live].astype(np.int64, copy=True)

    def result(self):
        if not self.done or not self.finished_flags[0]:raise ValueError('Beam search has no completed result')
        return self.finished[0, :1 + self.finished_lengths[0]].tolist(), float(self.finished_scores[0])