m1_compression/batched_arithmetic_coder.py

import torch
from typing import List, Tuple, Optional, Callable
from m1_compression import utils
import math

def _pdf_to_cdf(pdf: torch.Tensor) -> torch.Tensor:
    # NOTE: we do cumsum in float64, as we found
    # cumsum in float32 leads to numerical errors
    # when batch size is different across runs
    cdf = torch.cumsum(pdf.to(torch.float64), dim=-1).to(torch.float32)
    cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], dim=-1)  # prepend 0
    return cdf  # shape [..., vocab_size+1]

def _shift_left(x: int, base: int, base_to_pm1: int) -> int:
    """Shift `x` one digit left."""
    return (x % base_to_pm1) * base

def _shift_left_keeping_msd(x: int, base: int, base_to_pm1: int, base_to_pm2: int) -> int:
    """Shift `x` except MSD, which remains in place, one digit left."""
    return x - (x % base_to_pm1) + (x % base_to_pm2) * base

class BatchedArithmeticEncoder:
    def __init__(self, base: int, precision: int):
        self._base: int = base
        self._base_to_pm1: int = int(base ** (precision - 1))
        self._base_to_pm2: int = int(base ** (precision - 2))
        self._precision: int = precision

    def _get_bit_counts(self, buf_offsets: torch.Tensor, base_offsets: torch.Tensor) -> torch.Tensor:
        """Get bit counts for all sequences."""
        bit_counts = buf_offsets - base_offsets
        return bit_counts

    # Helper lambdas --------------------------------------------------------
    def flush_matching_digits(
            self, 
            low, 
            high, 
            old_low, 
            encoding: bool = True,
            bits_buffer: Optional[torch.Tensor] = None,
            buf_offsets: Optional[torch.Tensor] = None,
            num_carry_digits: Optional[torch.Tensor] = None,
            current_code_in_int: Optional[torch.Tensor] = None,
            _next_digit: Optional[Callable[[int], int]] = None,
            valid: Optional[torch.Tensor] = None, # add mask
        ):
        valid = valid if valid is not None else True
        while True:

            msd_low = low // self._base_to_pm1
            msd_high = high // self._base_to_pm1
            mask = msd_low == msd_high
            ## get masked mask
            mask = mask & valid
            if not torch.any(mask):
                break
     
            if encoding:
                msd_low_old = old_low // self._base_to_pm1
                # digit = msd_low[mask]
                # digit_old = msd_low_old[mask]
                # for idx, (d, d_old) in zip(mask.nonzero(as_tuple=False).flatten().tolist(), zip(digit.tolist(), digit_old.tolist())):
                #     # 1) real digit
                #     code[idx].append(int(d))
                #     # 2) any pending carries now resolved
                #     num_carry_digit = num_carry_digits[idx].item()
                #     if num_carry_digit:
                #         carry_digit = (
                #             self._base - 1 + d - d_old
                #         ) % self._base
                #         code[idx].extend([carry_digit] * num_carry_digit)
                #         num_carry_digits[idx] = 0

                # ------------------------------------------------------------------
                sel       = mask.nonzero(as_tuple=False).flatten()
                bits_buffer[buf_offsets[sel]] = msd_low[sel].to(torch.int32)
                buf_offsets.index_add_(
                    0, 
                    sel,
                    torch.ones_like(sel, dtype=torch.int32)
                )

                carry_sel = sel[(num_carry_digits[sel] > 0)]
                if carry_sel.numel():
                    _digit_carry = msd_low[carry_sel]
                    _digit_carry_old = msd_low_old[carry_sel]
                    carry_digit = (
                        self._base - 1 + _digit_carry - _digit_carry_old
                    ) % self._base

                    rep_cnt     = num_carry_digits[carry_sel]
                    repeats_max = rep_cnt.max().item()
                    grid        = torch.arange(
                        repeats_max,
                        device=rep_cnt.device
                    ).expand(carry_sel.size(0), repeats_max) # [K2, M]
                    mask_rep    = grid < rep_cnt.unsqueeze(1) # [K2, M]

                    start_pos   = buf_offsets[carry_sel]
                    target_pos  = (start_pos.unsqueeze(1) + grid)[mask_rep]
                    payload     = carry_digit.to(torch.int32).unsqueeze(1).expand_as(grid)[mask_rep]

                    bits_buffer[target_pos] = payload
                    buf_offsets.index_add_(0, carry_sel, rep_cnt)
                    num_carry_digits[carry_sel] = 0

            else:
                new_digit = torch.tensor([
                    _next_digit(i) if m else 0
                    for i, m in enumerate(mask.tolist())
                ], dtype=torch.int64, device=mask.device)
                current_code_in_int = torch.where(
                    mask, 
                    _shift_left(current_code_in_int, self._base, self._base_to_pm1) + new_digit, 
                    current_code_in_int
                )
            # Shift left to remove matching digits
            low = torch.where(
                mask, 
                _shift_left(low, self._base, self._base_to_pm1), 
                low
            )
            high = torch.where(
                mask,
                _shift_left(high, self._base, self._base_to_pm1) + self._base - 1,
                high,
            )
        return low, high, bits_buffer, buf_offsets, num_carry_digits, current_code_in_int

    def flush_carry_digits(
            self, 
            low, 
            high, 
            encoding: bool = True,
            num_carry_digits: Optional[torch.Tensor] = None,
            current_code_in_int: Optional[torch.Tensor] = None,
            _next_digit: Optional[Callable[[int], int]] = None,
            valid: Optional[torch.Tensor] = None,
    ):
        valid = valid if valid is not None else True
        while True:
            second_msd_low = (low // self._base_to_pm2) % self._base
            second_msd_high = ((high - 1) // self._base_to_pm2) % self._base
            msd_low = low // self._base_to_pm1
            msd_high = (high - 1) // self._base_to_pm1
            mask = (msd_low + 1 == msd_high) & (second_msd_low == self._base - 1) & (second_msd_high == 0)
            mask = mask & valid # valid mask
            if not torch.any(mask):
                break
            # For sequences in *mask*, we shift left *keeping MSD fixed*
            low = torch.where(
                mask, 
                _shift_left_keeping_msd(low, self._base, self._base_to_pm1, self._base_to_pm2), 
                low
            )
            high = torch.where(
                mask, 
                _shift_left_keeping_msd(high, self._base, self._base_to_pm1, self._base_to_pm2) + self._base - 1, 
                high
            )
            if encoding:
                num_carry_digits = torch.where(mask, num_carry_digits + 1, num_carry_digits)
            else:
                new_digit = torch.tensor([
                    _next_digit(i) if m else 0
                    for i, m in enumerate(mask.tolist())
                ], dtype=torch.int64, device=mask.device)
                current_code_in_int = torch.where(
                    mask, 
                    _shift_left_keeping_msd(current_code_in_int, self._base, self._base_to_pm1, self._base_to_pm2) + new_digit, 
                    current_code_in_int
                )
            
        return low, high, num_carry_digits, current_code_in_int

    def _process(
            self, 
            pdf: torch.Tensor,
            symbols: Optional[torch.Tensor] = None, 
            encoding: bool = True,
            return_num_padded_bits: bool = False,
            encoded_bits: Optional[List[List[int]]] = None,
            lengths: Optional[torch.Tensor] = None,
    ) -> List[bytes] | Tuple[List[bytes], List[int]]:
        assert pdf is not None, "symbols or pdf must be provided"
        assert pdf.ndim == 3, "input must be [B, T, V]"
        B, T, V = pdf.shape
        device = pdf.device
        if lengths is None:
            lengths = torch.full((B,),T,dtype = torch.int64,device=device)
        # --- interval state ----------------------------------------------------

        lengths = torch.clamp(lengths, min=0, max=T) 
        low = torch.zeros((B,), dtype=torch.int64, device=device)
        # NOTE: We represent the AC interval [0, 1) as rational numbers:
        #    [0, 1)
        #  ~ [self._low / base ** precision, (self._high + 1) / base ** precision)
        #  = [self._low / base ** precision, self._high / base ** precision],
        # where the we represent the upper bound *INCLUSIVE*. This is a subtle
        # detail required to make the integer arithmetic work correctly given that
        # all involved integers have `precision` digits in base `base`.
        high = torch.full((B,), int(self._base**self._precision) - 1, dtype=torch.int64, device=device)

        cdf = _pdf_to_cdf(pdf)

        if encoding:
            assert symbols is not None, "symbols must be provided for encoding"
            assert encoded_bits is None, "encoded_bits must be None for encoding"
            num_carry_digits = torch.zeros((B,), dtype=torch.int32, device=device)

            digits_sym  = math.ceil(math.log(V, self._base))
            max_digits  = self._precision + 2 * T * digits_sym
            bits_buffer = torch.empty(
                B * max_digits,
                dtype=torch.int32, 
                device=device
            )
            buf_offsets = torch.arange(
                B, 
                device=device, 
                dtype=torch.int32
            ) * max_digits
            _next_digit = None
            current_code_in_int = None
            decoded_symbols = None
        else:
            assert symbols is None, "symbols must be None for decoding"
            assert encoded_bits is not None, "encoded_bits must be provided for decoding"
            num_carry_digits = None
            bits_buffer = None
            buf_offsets = None
            # Stream read cursors -------------------------------------------------
            cursor = [0] * B
            def _next_digit(idx: int) -> int:
                if cursor[idx] < len(encoded_bits[idx]):
                    d = encoded_bits[idx][cursor[idx]]
                    cursor[idx] += 1
                    return d
                # Add padding to ensure the AC state is well-defined when decoding the last
                # symbol. Note that what exactly we do here depends on how encoder
                # termination is implemented (see `Encoder.terminate`).
                return self._base - 1

            current_code_in_int = torch.zeros((B,), dtype=torch.int64, device=device)
            for _ in range(self._precision):
                digits = torch.tensor([_next_digit(i) for i in range(B)],
                                    dtype=torch.int64, device=device)
                current_code_in_int = current_code_in_int * self._base + digits

            decoded_symbols = torch.zeros((B, T), dtype=torch.int64, device=device)
        # ---------------- main encoding ---------------------------------------
        for t in range(T):

            valid = t < lengths

            if not valid.any():
                break #all the sample is completed

            cdf_t = cdf[valid, t]  # [valid, V+1]
            low_valid = low[valid]
            high_valid = high[valid]
            width_valid = high_valid - low_valid + 1

            intervals = low_valid.unsqueeze(1) + (cdf_t * width_valid.unsqueeze(1)).type(torch.int64)

            if encoding:
                symbols_t = symbols[valid, t : t+1]
            else:
                symbols_t = torch.searchsorted(intervals, current_code_in_int[valid].unsqueeze(1), right=True) - 1
                # V 是词汇表大小，即 pdf.shape[-1]。
                symbols_t = symbols_t.clamp(max=V-1)
                # ==============================================================
                decoded_symbols[valid, t] = symbols_t.squeeze(1)
            
            old_low = low.clone()
            ## there is some wrong,if there are no valid sequences to process in a batch
            low[valid] = intervals.gather(1, symbols_t).squeeze(1)
            high[valid] = intervals.gather(1, (symbols_t + 1)).squeeze(1) - 1
            
            (low, high, bits_buffer, buf_offsets, num_carry_digits, current_code_in_int) = self.flush_matching_digits(
                low, high, old_low,
                encoding=encoding,
                bits_buffer=bits_buffer,
                buf_offsets=buf_offsets,
                num_carry_digits=num_carry_digits,
                current_code_in_int=current_code_in_int if not encoding else None,
                _next_digit=_next_digit if not encoding else None,
                valid=valid
            )

            (low, high, num_carry_digits, current_code_in_int) = self.flush_carry_digits(
                low, high,
                encoding=encoding,
                num_carry_digits=num_carry_digits,
                current_code_in_int=current_code_in_int if not encoding else None,
                _next_digit=_next_digit if not encoding else None,
                valid=valid
            )
            
        # for t in range(T):
        #     cdf_t = cdf[:, t]  # [B, V+1]
        #     width = high - low + 1
        #     intervals = low.unsqueeze(1) + (cdf_t * width.unsqueeze(1)).type(torch.int64)

        #     old_low = low

        #     if encoding:
        #         symbols_t = symbols[:, t:t+1]  # [B, 1]
        #     else:
        #         symbols_t = torch.searchsorted(intervals, current_code_in_int.unsqueeze(1), right=True) - 1
        #         decoded_symbols[:, t] = symbols_t.squeeze(1)

        #     low = intervals.gather(1, symbols_t).squeeze(1)
        #     high = intervals.gather(1, (symbols_t + 1)).squeeze(1) - 1

        #     # Renormalise until interval large enough
        #     (
        #         low, 
        #         high, 
        #         bits_buffer,         # encoding variable
        #         buf_offsets,         # encoding variable
        #         num_carry_digits,    # encoding variable
        #         current_code_in_int, # decoding variable
        #     ) = self.flush_matching_digits(
        #         low, 
        #         high, 
        #         old_low, 
        #         encoding, 
        #         bits_buffer,         # encoding variable
        #         buf_offsets,         # encoding variable
        #         num_carry_digits,    # encoding variable
        #         current_code_in_int, # decoding variable
        #         _next_digit,         # decoding variable
        #     )

        #     (
        #         low, 
        #         high, 
        #         num_carry_digits,    # encoding variable
        #         current_code_in_int, # decoding variable
        #     ) = self.flush_carry_digits(
        #         low, 
        #         high, 
        #         encoding,
        #         num_carry_digits,    # encoding variable
        #         current_code_in_int, # decoding variable
        #         _next_digit,         # decoding variable
        #     )

        if encoding:
            output_compressed_bytes = []
            output_num_padded_bits = []
            # -------------------- finalization ------------------------------------ 
            bits_buffer[buf_offsets] = (low // self._base_to_pm1).to(torch.int32)
            buf_offsets = buf_offsets + 1

            carry_sel = num_carry_digits.nonzero(as_tuple=False).flatten()
            if carry_sel.numel():
                carry_digit = self._base - 1

                rep_cnt     = num_carry_digits[carry_sel]
                repeats_max = rep_cnt.max()
                grid        = torch.arange(
                    repeats_max,
                    device=rep_cnt.device
                ).expand(carry_sel.size(0), repeats_max) # [K2, M]
                mask_rep    = grid < rep_cnt.unsqueeze(1) # [K2, M]

                start_pos   = buf_offsets[carry_sel]
                target_pos  = (start_pos.unsqueeze(1) + grid)[mask_rep]

                bits_buffer[target_pos] = carry_digit
                buf_offsets.index_add_(0, carry_sel, rep_cnt)
                num_carry_digits[carry_sel] = 0

            for idx in range(B):
                offset_start = idx * max_digits
                offset_end = buf_offsets[idx]
                bits_list = bits_buffer[offset_start:offset_end].cpu().tolist()
                bitstr = "".join(map(str, bits_list))
                compressed_bytes, num_padded_bits = utils.bits_to_bytes(bitstr)
                output_compressed_bytes.append(compressed_bytes)
                output_num_padded_bits.append(num_padded_bits)

            if return_num_padded_bits:
                return output_compressed_bytes, output_num_padded_bits
            else:
                return output_compressed_bytes
        else:
            return decoded_symbols

    def batched_encode(
            self, 
            pdf: torch.Tensor, 
            symbols: torch.Tensor,
            lengths: Optional[torch.Tensor] = None,
            return_num_padded_bits: bool = False
    ) -> List[bytes] | Tuple[List[bytes], List[int]]:
        B, T, V = pdf.shape
        device = pdf.device
        if lengths is None:
            lengths = torch.full((B,),T,dtype = torch.int64,device=device)
        return self._process(
            pdf, 
            symbols=symbols, 
            encoding=True, 
            return_num_padded_bits=return_num_padded_bits,
            encoded_bits=None, 
            lengths = lengths, ## pass length to address
        )

    def batched_decode(
            self, 
            pdf: torch.Tensor, 
            compressed_bytes: List[bytes], 
            num_padded_bits: List[int],
            lengths: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        B, T, V = pdf.shape
        device = pdf.device
        if lengths is None:
            lengths = torch.full((B,),T,dtype = torch.int64,device=device)
        assert len(compressed_bytes) == B, "encoded_bits length must be equal to batch size"
        assert len(num_padded_bits) == B, "num_padded_bits length must be equal to batch size"
        encoded_bits = [[] for _ in range(B)]
        for idx, (compressed_b, num_padded) in enumerate(zip(compressed_bytes, num_padded_bits)):
            bits = utils.bytes_to_bits(compressed_b, num_padded_bits=num_padded)
            encoded_bits[idx] = list(map(int, bits))
            # print("[DEBUG]: encoded_bits[{}]: {}".format(idx, encoded_bits[idx]))
        return self._process(
            pdf, 
            symbols=None, 
            encoding=False, 
            return_num_padded_bits=False,
            encoded_bits=encoded_bits, 
            lengths=lengths,
        )

    def incremental_batched_encode(
            self, 
            pdf: torch.Tensor, 
            symbols: torch.Tensor,
            lengths: Optional[torch.Tensor] = None,
            bit_threshold: Optional[int] = None,
            return_num_padded_bits: bool = False
    ) -> Tuple[List[bytes], List[int]] | Tuple[List[bytes], List[int], List[int]]:
        """
        Incrementally encode symbols with early stopping when bit threshold is exceeded.
        
        Args:
            pdf: [B, T, V] probability distributions
            symbols: [B, T] symbols to encode
            lengths: [B] length of each sequence (optional)
            bit_threshold: Stop encoding when any sequence exceeds this many bits
            return_num_padded_bits: Whether to return padding information
            
        Returns:
            final_compressed_bytes: List[bytes] - final compressed result for each sequence
            stopped_at_step: List[int] - step where each sequence stopped (-1 if completed normally)
            final_num_padded_bits: List[int] - padding info (only if return_num_padded_bits=True)
        """
        assert pdf.ndim == 3, "input must be [B, T, V]"
        B, T, V = pdf.shape
        device = pdf.device
        
        if lengths is None:
            lengths = torch.full((B,), T, dtype=torch.int64, device=device)
        
        lengths = torch.clamp(lengths, min=0, max=T)
        
        # Initialize arithmetic coding state
        low = torch.zeros((B,), dtype=torch.int64, device=device)
        high = torch.full((B,), int(self._base**self._precision) - 1, dtype=torch.int64, device=device)
        num_carry_digits = torch.zeros((B,), dtype=torch.int32, device=device)
        
        # Initialize bit buffer
        digits_sym = math.ceil(math.log(V, self._base))
        max_digits = self._precision + 2 * T * digits_sym
        bits_buffer = torch.empty(B * max_digits, dtype=torch.int32, device=device)
        buf_offsets = torch.arange(B, device=device, dtype=torch.int32) * max_digits

        base_offsets = torch.arange(B, device=device, dtype=torch.int32) * max_digits
        
        # Pre-allocate temporary buffers (avoid cloning at each step)
        temp_bits_buffer = torch.empty_like(bits_buffer)
        temp_buf_offsets = torch.empty_like(buf_offsets)
        temp_num_carry_digits = torch.empty_like(num_carry_digits)
        
        cdf = _pdf_to_cdf(pdf)
        
        # Track final results for each sequence - save buffer states, not bytes
        final_buffer = torch.empty_like(bits_buffer)
        final_buffer_ends = torch.zeros(B, dtype=torch.int32, device=device)
        final_num_padded_bits = [None] * B
        stopped_at_step = [-1] * B  # -1 means completed normally
        
        # Track which sequences are still active
        active_sequences = torch.ones(B, dtype=torch.bool, device=device)
        
        # Keep track of previous step's finalized buffer state for threshold logic
        prev_finalized_buffer = torch.empty_like(bits_buffer)
        prev_finalized_ends = torch.zeros_like(buf_offsets)
        
        for t in range(T):
            valid = (t < lengths) & active_sequences
            
            if not valid.any():
                break  # All sequences completed or stopped
            
            cdf_t = cdf[valid, t]  # [valid, V+1]
            low_valid = low[valid]
            high_valid = high[valid]
            width_valid = high_valid - low_valid + 1
            
            intervals = low_valid.unsqueeze(1) + (cdf_t * width_valid.unsqueeze(1)).type(torch.int64)
            symbols_t = symbols[valid, t:t+1]
            
            old_low = low.clone()
            low[valid] = intervals.gather(1, symbols_t).squeeze(1)
            high[valid] = intervals.gather(1, (symbols_t + 1)).squeeze(1) - 1
            
            # Flush digits and update buffers
            (low, high, bits_buffer, buf_offsets, num_carry_digits, _) = self.flush_matching_digits(
                low, high, old_low,
                encoding=True,
                bits_buffer=bits_buffer,
                buf_offsets=buf_offsets,
                num_carry_digits=num_carry_digits,
                current_code_in_int=None,
                _next_digit=None,
                valid=valid
            )
            
            (low, high, num_carry_digits, _) = self.flush_carry_digits(
                low, high,
                encoding=True,
                num_carry_digits=num_carry_digits,
                current_code_in_int=None,
                _next_digit=None,
                valid=valid
            )
            
            # Check if we need to compute results this step (if bit threshold checking or final step)
            need_check_threshold = bit_threshold is not None and active_sequences.any()
            some_seq_finished = ((t + 1 >= lengths) & active_sequences).any()
            
            if need_check_threshold or some_seq_finished:
                # Simulate finalization at this step using pre-allocated buffers
                temp_bits_buffer.copy_(bits_buffer, True)
                temp_buf_offsets.copy_(buf_offsets, True)
                temp_num_carry_digits.copy_(num_carry_digits, True)
                
                # Add final digit for all sequences (simulating termination)
                temp_bits_buffer[temp_buf_offsets] = (low // self._base_to_pm1).to(torch.int32)
                temp_buf_offsets += 1
                
                # Handle remaining carry digits for all sequences
                carry_sel = (temp_num_carry_digits > 0).nonzero(as_tuple=False).flatten()
                if carry_sel.numel():
                    carry_digit = self._base - 1
                    rep_cnt = temp_num_carry_digits[carry_sel]
                    repeats_max = rep_cnt.max()
                    grid = torch.arange(repeats_max, device=rep_cnt.device).expand(carry_sel.size(0), repeats_max)
                    mask_rep = grid < rep_cnt.unsqueeze(1)
                    
                    start_pos = temp_buf_offsets[carry_sel]
                    target_pos = (start_pos.unsqueeze(1) + grid)[mask_rep]
                    temp_bits_buffer[target_pos] = carry_digit
                    temp_buf_offsets.index_add_(0, carry_sel, rep_cnt)
                    temp_num_carry_digits[carry_sel] = 0
                
                # Check bit threshold and identify newly stopped sequences  
                if need_check_threshold:
                    current_bit_counts = self._get_bit_counts(temp_buf_offsets, base_offsets)
                    exceeds_threshold = (current_bit_counts > bit_threshold) & active_sequences
                    
                    if exceeds_threshold.any():
                        stopped_indices = exceeds_threshold.nonzero(as_tuple=False).flatten()
                        for idx in stopped_indices.cpu().tolist():  # Only move indices to CPU
                            active_sequences[idx] = False
                            stopped_at_step[idx] = t
                            # Save the result from PREVIOUS step (before exceeding threshold)
                            final_buffer_ends[idx] = prev_finalized_ends[idx]
                            offset_start = idx * max_digits
                            offset_end = prev_finalized_ends[idx]
                            final_buffer[offset_start:offset_end].copy_(prev_finalized_buffer[offset_start:offset_end])
                
                # If final step, all remaining active sequences need results
                is_final_step = (t + 1 >= lengths) & active_sequences
                if is_final_step.any():
                    final_step_indices = is_final_step.nonzero(as_tuple=False).flatten()
                    for idx in final_step_indices.cpu().tolist():
                        active_sequences[idx] = False
                        stopped_at_step[idx] = t + 1
                        # Save current step result for sequences that completed normally
                        final_buffer_ends[idx] = temp_buf_offsets[idx]
                        # Copy the finalized bits to main buffer for this sequence
                        offset_start = idx * max_digits
                        offset_end = temp_buf_offsets[idx]
                        final_buffer[offset_start:offset_end].copy_(temp_bits_buffer[offset_start:offset_end])
        
            # Update previous finalized buffer state for next iteration
            if need_check_threshold:
                prev_finalized_buffer.copy_(temp_bits_buffer)
                prev_finalized_ends.copy_(temp_buf_offsets)
        
        # Convert buffer states to compressed bytes at the very end
        final_compressed_bytes = []
        
        for idx in range(B):
            offset_start = idx * max_digits
            offset_end = final_buffer_ends[idx]
            bits_list = final_buffer[offset_start:offset_end].cpu().tolist()
            bitstr = "".join(map(str, bits_list))
            comp_bytes, num_padded = utils.bits_to_bytes(bitstr)
            final_compressed_bytes.append(comp_bytes)
            if return_num_padded_bits:
                final_num_padded_bits[idx] = num_padded
    
        if return_num_padded_bits:
            return final_compressed_bytes, stopped_at_step, final_num_padded_bits
        else:
            return final_compressed_bytes, stopped_at_step

def test_incremental_encoding():
    """Test the incremental encoding functionality"""
    print("Testing incremental encoding...")
    
    batch_size = 4
    seq_len = 32
    vocab_size = 64
    base = 2
    precision = 32
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Create test data with different scenarios
    symbols = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)
    pdf = torch.rand(batch_size, seq_len, vocab_size, device=device).clamp(min=1e-6)
    pdf = pdf.softmax(dim=-1)
    pdf = utils.batched_normalize_pdf_for_arithmetic_coding(pdf)
    # Test with variable lengths to ensure edge cases are covered
    lengths = torch.tensor([seq_len, seq_len-2, seq_len-4, seq_len], device=device)
    
    AC = BatchedArithmeticEncoder(base, precision)
    
    # Test incremental encoding with bit threshold
    bit_threshold = 20  # Stop when sequences exceed 20 bits
    final_compressed_bytes, stopped_at_step, final_num_padded_bits = AC.incremental_batched_encode(
        pdf, symbols, lengths=lengths, bit_threshold=bit_threshold, return_num_padded_bits=True
    )
    
    print(f"Stopped at steps: {stopped_at_step}")
    print(f"Final compressed sizes (bytes): {[len(cb) if cb else 0 for cb in final_compressed_bytes]}")
    bit_counts = [len(cb) * 8 - pb if cb else 0 for cb, pb in zip(final_compressed_bytes, final_num_padded_bits)]
    print(f"Final bit counts: {bit_counts}")
    
    # Test without threshold
    final_compressed_bytes_full, stopped_at_step_full = AC.incremental_batched_encode(
        pdf, symbols, lengths=lengths, bit_threshold=None, return_num_padded_bits=False
    )
    
    print(f"Full encoding stopped at: {stopped_at_step_full}")  # Should all be -1
    print(f"Full encoding sizes (bytes): {[len(cb) if cb else 0 for cb in final_compressed_bytes_full]}")
    
    # Consistency test: Verify that incremental encoding matches regular encoding on prefixes
    print("\n--- Consistency Test ---")
    
    # Extract prefixes based on where sequences stopped
    prefix_symbols = []
    prefix_pdf = []
    prefix_lengths = []
    
    for i in range(batch_size):
        if stopped_at_step[i] == -1:
            # Sequence completed normally, use full sequence
            stop_point = seq_len
        else:
            # Sequence stopped early, use up to stop point
            stop_point = stopped_at_step[i]
        
        prefix_symbols.append(symbols[i, :stop_point])
        prefix_pdf.append(pdf[i, :stop_point])
        prefix_lengths.append(stop_point)
    
    # Create tensors for prefix encoding
    max_prefix_len = max(prefix_lengths)
    batch_prefix_symbols = torch.zeros((batch_size, max_prefix_len), dtype=symbols.dtype, device=device)
    batch_prefix_pdf = torch.zeros((batch_size, max_prefix_len, vocab_size), dtype=pdf.dtype, device=device)
    
    for i in range(batch_size):
        length = prefix_lengths[i]
        batch_prefix_symbols[i, :length] = prefix_symbols[i]
        batch_prefix_pdf[i, :length] = prefix_pdf[i]
    
    prefix_lengths_tensor = torch.tensor(prefix_lengths, dtype=torch.int64, device=device)
    
    # Encode prefixes using regular batched encoder
    prefix_compressed_bytes, prefix_num_padded_bits = AC.batched_encode(
        batch_prefix_pdf, batch_prefix_symbols, 
        lengths=prefix_lengths_tensor, 
        return_num_padded_bits=True
    )
    
    # Compare results
    print("Comparing incremental vs regular encoding on prefixes:")
    all_match = True
    for i in range(batch_size):
        incremental_bytes = final_compressed_bytes[i]
        regular_bytes = prefix_compressed_bytes[i]
        
        if incremental_bytes == regular_bytes:
            print(f"Sequence {i}: ✓ Match (stopped at step {stopped_at_step[i]})")
        else:
            print(f"Sequence {i}: ✗ Mismatch (stopped at step {stopped_at_step[i]})")
            print(f"  Incremental: {len(incremental_bytes) if incremental_bytes else 0} bytes")
            print(f"  Regular:     {len(regular_bytes)} bytes")
            all_match = False
    assert all_match, "Some incremental encodings don't match - there may be a bug"
    if all_match:
        print("✓ All incremental encodings match regular encodings on prefixes!")
    else:
        print("✗ Some incremental encodings don't match - there may be a bug")
    
    print("✓ Incremental encoding test completed")

def test_consistency_various_thresholds():
    """Test consistency between incremental and regular encoding with various thresholds"""
    print("\n=== Testing Consistency with Various Thresholds ===")
    
    batch_size = 3
    seq_len = 128
    vocab_size = 4
    base = 2
    precision = 16
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Create fixed test data for reproducible results
    torch.manual_seed(42)
    symbols = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)
    pdf = torch.rand(batch_size, seq_len, vocab_size, device=device).clamp(min=1e-6)
    pdf = pdf.softmax(dim=-1)
    lengths = torch.full((batch_size,), seq_len, device=device)
    
    AC = BatchedArithmeticEncoder(base, precision)
    
    # Test with multiple threshold values
    thresholds = [10, 25, 50, 100, None]  # None means no threshold
    
    for threshold in thresholds:
        print(f"\n--- Testing with threshold: {threshold} ---")
        
        # Run incremental encoding
        if threshold is None:
            final_bytes, stop_steps = AC.incremental_batched_encode(
                pdf, symbols, lengths=lengths, bit_threshold=threshold
            )
        else:
            final_bytes, stop_steps = AC.incremental_batched_encode(
                pdf, symbols, lengths=lengths, bit_threshold=threshold
            )
        
        print(f"Stop steps: {stop_steps}")
        
        # Create prefix data based on stop points
        prefix_lengths = []
        for i in range(batch_size):
            if stop_steps[i] == -1:
                prefix_lengths.append(lengths[i].item())
            else:
                prefix_lengths.append(stop_steps[i])
        
        max_prefix_len = max(prefix_lengths)
        batch_prefix_symbols = torch.zeros((batch_size, max_prefix_len), dtype=symbols.dtype, device=device)
        batch_prefix_pdf = torch.zeros((batch_size, max_prefix_len, vocab_size), dtype=pdf.dtype, device=device)
        
        for i in range(batch_size):
            length = prefix_lengths[i]
            batch_prefix_symbols[i, :length] = symbols[i, :length]
            batch_prefix_pdf[i, :length] = pdf[i, :length]
        
        prefix_lengths_tensor = torch.tensor(prefix_lengths, dtype=torch.int64, device=device)
        
        # Run regular encoding on prefixes
        regular_bytes = AC.batched_encode(
            batch_prefix_pdf, batch_prefix_symbols, 
            lengths=prefix_lengths_tensor
        )
        
        # Compare results
        all_consistent = True
        for i in range(batch_size):
            if final_bytes[i] == regular_bytes[i]:
                print(f"  Seq {i}: ✓ (len={prefix_lengths[i]})")
            else:
                print(f"  Seq {i}: ✗ INCONSISTENT (len={prefix_lengths[i]})")
                print(f"    Incremental: {len(final_bytes[i]) if final_bytes[i] else 0} bytes")
                print(f"    Regular:     {len(regular_bytes[i])} bytes")
                all_consistent = False
        assert all_consistent, "Some incremental encodings don't match - there may be a bug"
        if all_consistent:
            print(f"  ✓ All sequences consistent for threshold {threshold}")
        else:
            print(f"  ✗ Inconsistencies found for threshold {threshold}")
    
    print("\n✓ Consistency testing completed")

if __name__ == "__main__":
    
    # Test incremental encoding
    test_incremental_encoding()
    
    # Test consistency between incremental and regular encoding
    test_consistency_various_thresholds()
    batch_size = 32
    seq_len = 128
    vocab_size = 256
    base = 2
    precision = 32
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    lengths = torch.randint(1,seq_len,(batch_size,),device=device)
    symbols = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)
    
    # 构造 one-hot pdf
    pdf = torch.zeros(batch_size, seq_len, vocab_size, device=device)
    for i in range(batch_size):
        # get pdf in valid length so it can be normalized.
        s = torch.randint(0, vocab_size, (lengths[i],), device=device)
        symbols[i, :lengths[i]] = s
        for t in range(lengths[i]):
            pdf[i, t, s[t]] = 1.0
    pdf = pdf / pdf.sum(-1, keepdim=True)

    pdf = torch.rand(batch_size, seq_len, vocab_size, device=device).clamp(min=1e-6)
    pdf = pdf.softmax(dim=-1)
    symbols = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)
    lengths = torch.randint(1, seq_len, (batch_size,), device=device)
    # lengths = torch.full((batch_size,), seq_len, device=device)

    AC = BatchedArithmeticEncoder(base, precision)

    # Test original functionality
    print("Testing original batched encoding...")
    start_event = torch.cuda.Event(enable_timing=True) if device == "cuda" else None
    end_event = torch.cuda.Event(enable_timing=True) if device == "cuda" else None
    
    if device == "cuda":
        start_event.record()
    codes, padded_bits = AC.batched_encode(pdf, symbols,lengths=lengths,return_num_padded_bits=True)
    if device == "cuda":
        end_event.record()
        torch.cuda.synchronize()
        print(f"CUDA wall clock time: {start_event.elapsed_time(end_event):.2f} ms")
    print([len(c) for c in codes])

    decoded = AC.batched_decode(pdf, codes, padded_bits,lengths)
    print("[DEBUG]: decoded {} symbols {}".format(decoded[0], symbols[0]))
    print("✓ passed - avg. digits per seq:", [len(s) for s in codes])

    # Validate the length
    for i in range(batch_size):
        print(lengths[i].item())
        print(decoded[i, :lengths[i].item()])
        print(symbols[i, :lengths[i].item()])
        l = lengths[i].item()
        assert torch.all(decoded[i, :l] == symbols[i, :l]), f"Sample {i} mismatch"
    print("All variable-length sequences verified successfully")