import torch
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence

texts = [
    [5, 9, 12, 34, 23, 76, 89],          # len = 7
    [2, 45, 23],                         # len = 3
    [12, 56, 78, 90, 43, 22, 11, 8],     # len = 8
    [65, 12, 99, 54],                    # len = 4
    [3, 7, 8, 5, 1, 4, 6, 10, 11],       # len = 9
]

labels = [0, 1, 0, 2, 1]

class TextDataset(Dataset):
    def __init__(self, texts, labels):
        # сложите тексты и классы в переменные датасета
        self.texts = texts
        self.labels = labels

    def __len__(self):
        # посчитайте размер датасета
        return len(self.texts)

    def __getitem__(self, idx):
        # верните элемент датасета
        return {
            'text': torch.tensor(self.texts[idx], dtype=torch.long),
            'label': torch.tensor(self.labels[idx], dtype=torch.long)
        }

def collate_fn(batch):
    texts = [item['text'] for item in batch]
    labels = [item['label'] for item in batch]

    # посчитайте длины текстов в батче
    lengths = []
    for _ in range(len(texts)): 
        lengths.append(len(texts[_]))

    # отсортируйте тексты и классы по убыванию длины текстов
    joined_data = list(zip(lengths, texts, labels))
    # joined_data = torch.cat(
    #     torch.tensor(lengths), 
    #     torch.tensor(texts), 
    #     labels
    #     )
    joined_data.sort(key=lambda x: x[0], reverse=True)
    # joined_data = torch.argsort(joined_data, dim=0)
    lengths, texts, labels = zip(*joined_data)
    # lengths, texts, labels = torch.split(joined_data, 0)

    # дополните тексты пэддингом
    padded_texts = pad_sequence(texts, batch_first=True, padding_value=0)

    # посчитайте маску для батча
    masks = (padded_texts != 0).long()

    return {
        'texts': padded_texts,     # (batch_size, max_len)
        'masks': masks,            # (batch_size, max_len)
        'labels': labels,          # (batch_size,)
        'lengths': lengths
    }

dataset = TextDataset(texts, labels)

# создайте объект dataloader с batch_size=3 и реализованным collate_fn
dataloader = DataLoader(dataset, batch_size=3, collate_fn=collate_fn)

# вывод результатов
for batch in dataloader:
    print("Texts:\n", batch['texts'])
    print("Masks:\n", batch['masks'])
    print("Labels:", batch['labels'])
    # print("Lengths:", batch['lengths'])
    print("=" * 40)