import torch
from torch import Tensor


def columnwise_clamp(
    X: Tensor, lower: Tensor, upper: Tensor, raise_on_violation=False
) -> Tensor:
    """Clamp values of a ndarray in column-wise fashion (with support for t-batches).

    Args:
        X: The 'b x n x d' input tensor. If 2-dimensional, 'b' is assumed to be 1.
        lower: The column-wise lower bounds. If scalar, apply bound to all columns.
        upper: The column-wise upper bounds. If scalar, apply bound to all columns.
        raise_on_violation: If `True`, raise an exception when the elments in `X`
            are out of the specified bounds (up to numerical accuracy). This is
            useful for post-processing candidates generated by optimizers that
            satisfy imposed bounds only up to numerical accuracy.

    Returns:
        The clamped ndarray.
    """
    min_bounds = lower.expand_as(X)  # np.broadcast_to(lower, X.shape)
    max_bounds = upper.expand_as(X)  # np.broadcast_to(upper, X.shape)
    if torch.any(min_bounds > max_bounds):
        raise ValueError("Minimum values must be <= maximum values")
    Xout = X.max(min_bounds)  # np.maximum(X, min_bounds)
    Xout = Xout.min(max_bounds)  # np.minimum(Xout, max_bounds)
    if raise_on_violation and not torch.allclose(Xout, X):
        raise RuntimeError("Original value(s) are out of bounds.")
    return Xout
