<gh_stars>0
import sys

sys.path.extend(['../'])

import pickle
import numpy as np
from torch.utils.data import Dataset

from feeders import tools


class Feeder(Dataset):
    def __init__(self, data_path, label_path, seq_len_path, using_config=None, mode='train'):
        self.data_path = data_path
        self.label_path = label_path
        self.seq_len_path = seq_len_path
        self.time_patch_size = using_config.model.window_size // using_config.model.seq_num
        self.using_config = using_config
        self.augmentation = not using_config.augmentation.disable_all_aug
        self.mode = mode
        self.clip_range = float(using_config.augmentation.clip_range)
        if self.using_config.augmentation.aug_resize_clip is not True:
            self.clip_range = 0

        if self.mode == 'test':
            self.augmentation = self.using_config.augmentation.test_aug
            self.clip_range = 0

        self.load_data()
        pass

    def load_data(self):
        try:
            with open(self.label_path) as f:
                self.label = pickle.load(f)
        except:
            # for pickle file from python2
            with open(self.label_path, 'rb') as f:
                self.label = pickle.load(f, encoding='latin1')

        try:
            with open(self.data_path) as f:
                self.data = pickle.load(f)
        except:
            # for pickle file from python2
            with open(self.data_path, 'rb') as f:
                self.data = pickle.load(f, encoding='latin1')

        try:
            with open(self.seq_len_path) as f:
                self.seq_len = pickle.load(f)
        except:
            # for pickle file from python2
            with open(self.seq_len_path, 'rb') as f:
                self.seq_len = pickle.load(f, encoding='latin1')

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

    def __iter__(self):
        return self

    def __getitem__(self, index):
        data_numpy = self.data[index]
        label = self.label[index]
        seq_len = self.seq_len[index]
        data_numpy = np.array(data_numpy)

        # temporal aug
        data_numpy = tools.random_resize_and_clip(data_numpy, clip_range=self.clip_range,
                                                  size=self.using_config.model.window_size)

        # spatial aug
        if self.augmentation:
            T, V, C = data_numpy.shape
            data_numpy = np.ascontiguousarray(data_numpy)
            if self.using_config.augmentation.patch_level_augument:
                F = self.using_config.model.seq_num
                data_list = data_numpy.reshape((F, self.time_patch_size, V, C))
            else:
                data_list = np.expand_dims(data_numpy, axis=0)

            data_list = np.ascontiguousarray(data_list)
            for i in range(data_list.shape[0]):
                if self.using_config.augmentation.aug_shear:
                    data_list[i] = tools.shear(data_list[i])

                if self.using_config.augmentation.aug_zoom:
                    data_list[i] = tools.randomZoom(data_list[i])

                if self.using_config.augmentation.aug_rotate:
                    data_list[i] = tools.randomRotate(data_list[i])

                if self.using_config.augmentation.aug_spatial_flip:
                    data_list[i] = tools.randomSpatialFlip(data_list[i])

                if self.using_config.augmentation.aug_sptial_shift:
                    data_list[i] = tools.randomSpatialShift(data_list[i])

                if self.using_config.augmentation.aug_joint_courruption:
                    data_list[i] = tools.joint_courruption(data_list[i])

            data_list = np.ascontiguousarray(data_list)
            data_numpy = data_list.reshape(T, V, C)
        return data_numpy, label, seq_len, index

    def top_k(self, score, top_k):
        rank = score.argsort()
        hit_top_k = [l in rank[i, -top_k:] for i, l in enumerate(self.label)]
        return sum(hit_top_k) * 1.0 / len(hit_top_k)
