import torch

from torch.utils.data import Dataset, DataLoader

from torch.utils.data.sampler import Sampler, RandomSampler

"""

class TmpDataset(Dataset):

    def __init__(self, m=10):

        self.len = m

    def __getitem__(self, index):

        return (list(range(10)) * index, [0] * index)

    def __len__(self):

        return self.len

"""

class FixedLengthBatchSampler(Sampler):

    def __init__(self, sampler, fixed_length, drop_last):

        self.sampler = sampler

        self.fixed_length = fixed_length

        self.drop_last = drop_last

        self.rel_sampler_count = 0

    def __iter__(self):

        batch = []

        now_length = 0

        for idx in self.sampler:

            #print(batch, now_length)

            sample_length = len(self.sampler.data_source[idx][-1]) * 3

            if now_length + sample_length > self.fixed_length:

                #print(batch, now_length)

                yield batch

                batch = []

                now_length = 0

            batch.append(idx)

            now_length += sample_length

            self.rel_sampler_count += 1

        if len(batch) > 0 and not self.drop_last:

            yield batch

def my_collate_fn(batch):

    type_count = len(batch[0])

    batch_size = sum([len(item[-1]) for item in batch])

    output = ()

    for i in range(type_count):

        tmp = []

        for item in batch:

            tmp.extend(item[i])

        if len(tmp) <= batch_size:

            output += (torch.LongTensor(tmp),)

        else:

            output += (torch.LongTensor(tmp).reshape(batch_size, -1),)

    return output