import torch
import numpy as np
import random

from collections import OrderedDict
import os
import glob
import cv2
import torch.utils.data as data


def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    # torch.backends.cudnn.deterministic = True # 加了会变慢
    return seed


def np_load_frame(filename, resize_height, resize_width):
    """
    Load image path and convert it to numpy.ndarray. Notes that the color channels are BGR and the color space
    is normalized from [0, 255] to [-1, 1].

    :param filename: the full path of image
    :param resize_height: resized height
    :param resize_width: resized width
    :return: numpy.ndarray
    """
    image_decoded = cv2.imread(filename)
    image_resized = cv2.resize(image_decoded, (resize_width, resize_height))
    image_resized = image_resized.astype(dtype=np.float32)
    image_resized = (image_resized / 127.5) - 1.0
    return image_resized


class DataLoader(data.Dataset):
    def __init__(
        self,
        args,
        video_folder,
        transform,
        resize_height,
        resize_width,
        time_step=4,
        num_pred=1,
        is_concat=False,
    ):
        self.args = args
        self.dir = video_folder  # video_folder
        self.transform = transform
        self.videos = OrderedDict()
        self._resize_height = resize_height
        self._resize_width = resize_width
        self._time_step = time_step
        self._num_pred = num_pred
        self.setup()
        self.samples = self.get_all_samples()
        self.is_concat = is_concat

    def setup(self):
        videos = glob.glob(os.path.join(self.dir, "*"))
        videos2 = glob.glob(
            os.path.join(self.dir.replace("datasets", "datasets2"), "*")
        )

        # videos.extend(videos2)
        clip_path = self.dir.replace(self.args.dataset_path, self.args.clip_feas_path)

        # print('dir:', self.dir)
        # print(self.args.dataset_path, self.args.clip_feas_path)
        # print(clip_path)
        # aaaa

        for video in sorted(videos):
            video_name = video.split("/")[-1]
            self.videos[video_name] = {}
            self.videos[video_name]["path"] = video
            self.videos[video_name]["frame"] = glob.glob(os.path.join(video, "*.jpg"))
            self.videos[video_name]["frame"].sort()
            self.videos[video_name]["length"] = len(self.videos[video_name]["frame"])

            self.videos[video_name]["clip"] = np.load(
                os.path.join(clip_path, "feas_{}.npy".format(video_name))
            )

    def get_all_samples(self):
        frames = []
        videos = glob.glob(os.path.join(self.dir, "*"))
        for video in sorted(videos):
            video_name = video.split("/")[-1]
            for i in range(len(self.videos[video_name]["frame"]) - self._time_step):
                frames.append(self.videos[video_name]["frame"][i])

        return frames

    def __getitem__(self, index):
        video_name = self.samples[index].split("/")[-2]
        frame_name = int(self.samples[index].split("/")[-1].split(".")[-2])

        batch = []
        feas_batch = []
        for i in range(self._time_step + self._num_pred):
            image = np_load_frame(
                self.videos[video_name]["frame"][frame_name + i],
                self._resize_height,
                self._resize_width,
            )

            feas = self.videos[video_name]["clip"][frame_name + i]
            feas_batch.append(feas)

            if self.transform is not None:
                batch.append(self.transform(image))
        if self.is_concat:
            return [np.concatenate(batch, axis=0), np.concatenate(feas_batch, axis=0)]
        else:
            return [np.array(batch), np.array(feas_batch, dtype=np.float32)]

    def __len__(self):
        return len(self.samples)
