"""Padding utilities — pad variable-length sequences to a common length."""

from __future__ import annotations

import numpy as np


def pad_sequences(
    sequences: list[np.ndarray],
    max_length: int | None = None,
    pad_value: float = 0.0,
    dtype: np.dtype = np.float32,
) -> tuple[np.ndarray, np.ndarray]:
    """Pad a list of 2D sequences to a common length.

    Args:
        sequences: List of (T_i × F) arrays.
        max_length: Target sequence length. If None, uses max observed.
        pad_value: Fill value for padded positions.
        dtype: Output dtype.

    Returns:
        (padded_values, padding_mask) where:
            padded_values: (S × max_length × F) padded tensor.
            padding_mask: (S × max_length) boolean mask (True = padded).
    """
    if not sequences:
        return np.array([], dtype=dtype), np.array([], dtype=bool)

    S = len(sequences)
    F = sequences[0].shape[1]
    max_len = max_length or max(s.shape[0] for s in sequences)

    padded = np.full((S, max_len, F), pad_value, dtype=dtype)
    mask = np.ones((S, max_len), dtype=bool)

    for i, seq in enumerate(sequences):
        length = min(seq.shape[0], max_len)
        padded[i, :length, :] = seq[:length, :]
        mask[i, :length] = False

    return padded, mask


def unpad_sequences(
    padded_values: np.ndarray,
    padding_mask: np.ndarray,
) -> list[np.ndarray]:
    """Extract ragged sequences from a padded tensor.

    Args:
        padded_values: (S × T_max × F) padded tensor.
        padding_mask: (S × T_max) boolean mask (True = padded).

    Returns:
        List of (T_i × F) arrays.
    """
    sequences = []
    for s in range(padded_values.shape[0]):
        valid = ~padding_mask[s]
        sequences.append(padded_values[s, valid, :])
    return sequences