"""
Modified based on: https://github.com/open-mmlab/mmskeleton
open-mmlab/mmskeleton is licensed under the Apache License 2.0
"""

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


class Feeder(Dataset):
    def __init__(self, data_path, label_path, normalization=False, use_mmap=True):
        """
        :param data_path:
        :param label_path:
        :param normalization: If true, normalize input sequence
        :param use_mmap: If true, use mmap mode to load data, which can save the running memory
        """
        self.data_path = data_path
        self.label_path = label_path
        self.normalization = normalization
        self.use_mmap = use_mmap
        self.load_data()
        if normalization:
            self.get_mean_map()

    def load_data(self):
        # data: N C V T M
        try:
            with open(self.label_path) as f:
                self.sample_name, self.label = pickle.load(f)
        except Exception:
            # for pickle file from python2
            with open(self.label_path, 'rb') as f:
                self.sample_name, self.label = pickle.load(f, encoding='latin1')

        # load data
        if self.use_mmap:
            self.data = np.load(self.data_path, mmap_mode='r')
        else:
            self.data = np.load(self.data_path)

    def get_mean_map(self):
        data = self.data
        N, C, T, V, M = data.shape
        self.mean_map = data.mean(axis=2, keepdims=True).mean(axis=4, keepdims=True).mean(axis=0)
        self.std_map = data.transpose((0, 2, 4, 1, 3)).reshape((N * T * M, C * V)).std(axis=0).reshape((C, 1, V, 1))

    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]
        data_numpy = np.array(data_numpy)
        if self.normalization:
            data_numpy = (data_numpy - self.mean_map) / self.std_map

        return data_numpy, label, 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)
