# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/models/torch/utils.ipynb.

# %% auto 0
__all__ = ['get_padded_sequence']

# %% ../../../nbs/models/torch/utils.ipynb 3
import torch
import torch.nn.functional as F

# %% ../../../nbs/models/torch/utils.ipynb 6
def get_padded_sequence(data, sequence_length, idx, value=0):
    if not isinstance(data, torch.Tensor):
        data = torch.tensor(data)
    
    assert sequence_length % 2 == 1
    assert idx >= 0
    assert idx < data.size(0)

    data_length = data.size(0)
    start_idx = idx - sequence_length // 2
    end_idx = idx + sequence_length // 2 + 1

    # Pad the sequence with zeros if it's shorter than sequence_length
    if start_idx < 0:
        sequence = data[:end_idx]
        return F.pad(sequence, (-(start_idx), 0), mode='constant', value=value)
    elif end_idx > data_length:
        sequence = data[start_idx:]
        return F.pad(sequence, (0, end_idx - data_length), mode='constant', value=value)


    return data[start_idx:end_idx]
