import h5py
import hdf5storage
import numpy as np

SUMME_GT_DIR = '../../data/SumMe/GT'
PROCESSED_SUMME = '../../data/SumMe/processed/eccv16_dataset_summe_google_pool5.h5'


# PROCESSED_TVSUM = '../../data/TVSUM/processed/eccv16_dataset_tvsum_google_pool5.h5'


def load_tvsum_mat(filename):
    data = hdf5storage.loadmat(filename, variable_names=['tvsum50'])
    data = data['tvsum50'].ravel()

    data_list = []
    for item in data:
        video, category, title, length, nframes, user_anno, gt_score = item

        item_dict = {
            'video': video[0, 0],
            'category': category[0, 0],
            'title': title[0, 0],
            'length': length[0, 0],
            'nframes': nframes[0, 0],
            'user_anno': user_anno,
            'gt_score': gt_score
        }

        data_list.append(item_dict)

    return data_list


def load_processed_summe(processed_summe=PROCESSED_SUMME):
    with h5py.File(processed_summe, 'r') as hdf:
        data_list = dict()
        for video_name in hdf.keys():
            video_idx = video_name[6:]
            elemnt = 'video_' + video_idx
            nframes = np.array(hdf.get(elemnt + '/n_frames'))
            positions = np.array(hdf.get(elemnt + '/picks'))
            user_score = np.array(hdf.get(elemnt + '/user_summary'))
            mean_user_score = np.asarray(user_score).mean(axis=0)
            gt_score = np.array(hdf.get(elemnt + '/gtsummary'))
            item_dict = {
                'nframes': nframes,
                'user_score': user_score,
                'positions': positions,
                'mean_user_score': mean_user_score,
                'gt_score': gt_score,
            }

            data_list[elemnt] = item_dict

    hdf.close()
    return data_list