from __future__ import annotations

from typing import Any, Iterable, Sequence

import numpy as np
import torch


def pad_sequences(
    sequences: Iterable[Any],
    *,
    maxlen: int | None = None,
    dtype: str | np.dtype = np.float32,
    padding: str = "post",
    truncating: str = "post",
    value: float = 0.0,
    return_lengths: bool = True,
) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
    """Pad a list of variable-length sequences to a uniform length.

    Parameters mirror the common Keras-style `pad_sequences` behavior, but support
    multi-feature sequences with shape ``(T, F...)``.
    """
    seq_list = [np.asarray(seq) for seq in sequences]
    if not seq_list:
        raise ValueError("sequences must contain at least one item.")
    if padding not in {"pre", "post"}:
        raise ValueError("padding must be 'pre' or 'post'.")
    if truncating not in {"pre", "post"}:
        raise ValueError("truncating must be 'pre' or 'post'.")

    for idx, seq in enumerate(seq_list):
        if seq.ndim == 0:
            raise ValueError(f"Sequence at index {idx} must have at least 1 dimension (time, ...).")

    feature_shape = tuple(seq_list[0].shape[1:])
    for idx, seq in enumerate(seq_list[1:], start=1):
        if tuple(seq.shape[1:]) != feature_shape:
            raise ValueError(
                f"Inconsistent feature shape at index {idx}: expected {feature_shape!r}, got {tuple(seq.shape[1:])!r}."
            )

    raw_lengths = np.asarray([int(seq.shape[0]) for seq in seq_list], dtype=np.int64)
    if np.any(raw_lengths < 0):
        raise ValueError("Sequence lengths must be non-negative.")
    inferred_maxlen = int(raw_lengths.max()) if raw_lengths.size else 0
    if maxlen is None:
        maxlen = inferred_maxlen
    maxlen = int(maxlen)
    if maxlen < 0:
        raise ValueError("maxlen must be >= 0.")

    out_shape = (len(seq_list), maxlen, *feature_shape)
    padded = np.full(out_shape, value, dtype=np.dtype(dtype))
    clipped_lengths = np.minimum(raw_lengths, maxlen)

    for i, seq in enumerate(seq_list):
        if maxlen == 0:
            continue
        if truncating == "pre":
            trunc = seq[-maxlen:]
        else:
            trunc = seq[:maxlen]
        n = int(trunc.shape[0])
        if n == 0:
            continue
        if padding == "pre":
            padded[i, -n:] = trunc
        else:
            padded[i, :n] = trunc

    if return_lengths:
        return padded, clipped_lengths
    return padded


def sequence_collate_fn(batch: Sequence[Any], *, pad_value: float = 0.0) -> Any:
    """Collate variable-length sequences for a PyTorch DataLoader.

    Supported batch formats:
    - ``[sequence, ...]`` -> ``(x_padded, lengths)``
    - ``[(sequence, target), ...]`` -> ``(x_padded, targets, lengths)``
    """
    if len(batch) == 0:
        raise ValueError("batch must contain at least one item.")

    first = batch[0]
    if isinstance(first, (tuple, list)):
        if len(first) != 2:
            raise ValueError("sequence_collate_fn currently supports batch items of shape (sequence, target).")
        seqs = [item[0] for item in batch]
        targets = [item[1] for item in batch]
        x_padded, lengths = pad_sequences(seqs, value=pad_value, return_lengths=True)
        x_tensor = torch.as_tensor(x_padded, dtype=torch.float32)
        y_tensor = torch.as_tensor(np.asarray(targets))
        lengths_tensor = torch.as_tensor(lengths, dtype=torch.long)
        return x_tensor, y_tensor, lengths_tensor

    x_padded, lengths = pad_sequences(batch, value=pad_value, return_lengths=True)
    x_tensor = torch.as_tensor(x_padded, dtype=torch.float32)
    lengths_tensor = torch.as_tensor(lengths, dtype=torch.long)
    return x_tensor, lengths_tensor

