import functools
from io import BytesIO

import torch
import torchvision
import torchvision.transforms.v2 as transforms
import wids
from torch.utils.data import DataLoader


def _video_shortener(video_tensor, length, generator=None):
    start = torch.randint(0, video_tensor.shape[0] - length, (1,), generator=generator)
    return video_tensor[start:start + length]


def select_video_extract(length=16, generator=None):
    return functools.partial(_video_shortener, length=length, generator=generator)


def my_collate_fn(batch):
    videos = torch.stack([sample[0] for sample in batch])
    txts = [sample[1] for sample in batch]

    return videos, txts


class WebVidDataset(wids.ShardListDataset):

    def __init__(self, shards, cache_dir, video_length=16, video_size=256, video_length_offset=1, val=False, seed=42,
                 **kwargs):

        self.val = val
        self.generator = torch.Generator()
        self.generator.manual_seed(seed)
        self.generator_init_state = self.generator.get_state()
        super().__init__(shards, cache_dir=cache_dir, keep=True, **kwargs)

        if isinstance(video_size, int):
            video_size = (video_size, video_size)

        self.video_size = video_size

        for size in video_size:
            if size % 8 != 0:
                raise ValueError("video_size must be divisible by 8")

        self.transform = transforms.Compose(
            [
                select_video_extract(length=video_length + video_length_offset, generator=self.generator),
                transforms.Resize(size=video_size),
                transforms.RandomCrop(size=video_size) if not self.val else transforms.CenterCrop(size=video_size),
                transforms.RandomHorizontalFlip() if not self.val else transforms.Identity(),
            ]
        )

        self.add_transform(self._make_sample)

    def _make_sample(self, sample):
        if self.val:
            self.generator.set_state(self.generator_init_state)
        video = torchvision.io.read_video(BytesIO(sample[".mp4"].read()), output_format="TCHW", pts_unit='sec')[0]
        label = sample[".txt"]
        return self.transform(video), label


if __name__ == "__main__":

    dataset = WebVidDataset(
        tar_index=0,
        root_path='/users/Etu9/3711799/onlyflow/data/webvid/desc.json',
        video_length=16,
        video_size=256,
        video_length_offset=0,
    )

    sampler = wids.DistributedChunkedSampler(dataset, chunksize=1000, shuffle=True)
    dataloader = DataLoader(
        dataset,
        collate_fn=my_collate_fn,
        batch_size=4,
        sampler=sampler,
        num_workers=4
    )

    for i, (images, labels) in enumerate(dataloader):
        print(i, images.shape, labels)
        if i > 10:
            break