from __future__ import print_function
import torch.utils.data as data
import os
import sys
import numpy as np
import h5py

class MLMLoader(data.Dataset):
    def __init__(self, data_path, partition, mismatch=0.5):

        if data_path == None:
            raise Exception('No data path specified.')

        if partition is None or partition not in ['train', 'val', 'test']:
            raise Exception('Unknown partition type %s.' % partition)
        else:
            self.partition = partition

        self.h5f = h5py.File(os.path.join(data_path, f'{partition}.h5'), 'r')

        self.ids = self.h5f['ids']
        self.mismatch = mismatch

    def __getitem__(self, index):
        instanceId = self.ids[index]
        # we force 50 percent of them to be a mismatch
        match = np.random.uniform() > self.mismatch if self.partition == 'train' else True

        target = match and 1 or -1

        if target == 1:  # load positive samples

            all_img = self.h5f[f'{instanceId}_images'][()]
            if self.partition == 'train':
                img = all_img[np.random.choice(range(all_img.shape[0]))]
            else:
                img = all_img[0]

            summaries = self.h5f[f'{instanceId}_summaries'][()]
            multi_wiki = summaries[np.random.choice(range(np.shape(summaries)[0]))]
            triple = self.h5f[f'{instanceId}_classes'][()]
            triple = triple[np.random.choice(range(triple.shape[0]))]
            cluster = self.h5f[f'{instanceId}_onehot'][()]

        else:
            # Negative samples are generated by picking random images
            # Other modalities remain unchanged
            all_idx = range(len(self.ids))
            rndindex = np.random.choice(all_idx)
            # random index to pick image ids at random
            while rndindex == index:
                rndindex = np.random.choice(all_idx)  # pick a random index

            # load negative samples of images
            rndId = self.ids[rndindex]
            all_img = self.h5f[f'{rndId}_images']

            # negative samples are in train only
            if self.partition == 'train':
                img = all_img[np.random.choice(range(all_img.shape[0]))]
            else:
                img = all_img[0]
            
            # other modalities remain unchanged
            summaries = self.h5f[f'{rndId}_summaries'][()]
            multi_wiki = summaries[np.random.choice(range(np.shape(summaries)[0]))]
            triple = self.h5f[f'{instanceId}_classes'][()]
            triple = triple[np.random.choice(range(triple.shape[0]))]
            cluster = self.h5f[f'{instanceId}_onehot'][()]

        # output
        output = {
            'image': img,
            'multi_wiki': multi_wiki,
            'target': target,
            'cluster': cluster,
            'triple': triple
        }

        if self.partition != 'train':
            output['id'] = instanceId

        return output

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