import torch


def become_constant_trick(x, lengths):
    """
    Input:
        x: a tensor of shape (batch, stream, channel)
        lengths: a tensor of shape (batch,)

    Returns:
        A tensor y of shape (batch, stream, channel), such that
        y[i, j, k] = x[i, j, k] if j < lengths[i] else x[i, lengths[i], k]

    This is useful when taking the signature afterwards, as it makes the signature not notice any changes made after
    length.
    """

    batch, stream, channel = x.shape
    if (stream < lengths).any():
        raise ValueError("x's stream dimension is of length {} but one of the lengths is longer than this. lengths={}"
                         "".format(stream, lengths))
    lengths = lengths - 1  # go from length-of-sequence to index-of-final-element-in-sequence
    expanded_lengths = lengths.unsqueeze(1).unsqueeze(2).expand(batch, 1, channel)
    final_value = x.gather(dim=1, index=expanded_lengths)
    final_value.expand(batch, stream, channel)
    mask = torch.arange(0, stream, device=x.device).unsqueeze(0) > lengths.unsqueeze(1)
    mask = mask.unsqueeze(2).expand(batch, stream, channel)
    return x.masked_scatter(mask, final_value.masked_select(mask))