from __future__ import print_function
import tensorflow as tf
import numpy as np
from PIL import Image
import os
import glob
import csv
import time

class dataset_api(object):
    def __init__(self, n_examples, im_height, im_width, channels, n_episodes, mean_sub):
        self.n_examples = n_examples
        self.im_height = im_height
        self.im_width = im_width
        self.channels = channels
        self.n_episodes = n_episodes
        self.root_dir = '../data/miniImagenet'

        self.mean_sub = mean_sub # if use mean substraction for preprocessing

        self.mean_value = [120.15719937, 114.71930599, 102.78186757]

    def subtract_mean(self, X):
        # per image mean subtraction
        # X: N_c x N x H x W x C
        '''
        N_c,N,H,W,C = np.shape(X)
        Xf = np.reshape(X, [-1,C])
        means = np.mean(Xf, axis=0, dtype=np.float64) # float64 is necessray
        print(means)
        '''

        return X - self.mean_value


    def load_data(self, n_way, n_shot, n_query, stage='train', n_epochs=1000):
        """"
            Load data main func: use stage params to decide train, val or test
        """
         ##  main dataset code from here
        classes = os.listdir(os.path.join(self.root_dir,'data', stage))
        n_classes = len(classes)
        print('Load {}, {} classes'.format(stage, n_classes))

        all_eps = np.zeros((n_epochs*100, n_way), dtype=np.int32)
        for itr in range(n_epochs*100):
            epi_classes = np.random.permutation(n_classes)[:n_way]
            all_eps[itr] = epi_classes

        dataset = tf.data.Dataset.from_tensor_slices(all_eps)

        def _parse_py_function(ep_class):
            """ pyfunc to deal with dirs and img names """
            selected = [[]]*len(ep_class)
            for i in range(len(ep_class)):
                cls_name   = classes[ep_class[i]]
                cls_pattern = self.root_dir + '/data/'+stage+'/' + cls_name +'/*jpg'
                imgs       = glob.glob(cls_pattern)

                np.random.shuffle(imgs)  # in-place operation
                selected[i] = imgs[0:n_shot+n_query]

            selected = np.array(selected)
            return selected

        def precess_imgs(selected):
            """ tf operation to load and preprocess images """
            selected = tf.reshape(selected, [n_way,n_shot+n_query])
            out_imgs = []

            for i in range(n_way):
                out_imgs.append([])
                for j in range(n_shot+n_query):
                    fname = selected[i][j]
                    image_string = tf.read_file(fname)
                    image_decoded = tf.image.decode_jpeg(image_string)
                    image_resized = tf.image.resize_images(image_decoded, [self.im_height, self.im_width],
                                        method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
                    if self.mean_sub:
                        image_resized = self.subtract_mean(image_resized)
                    out_imgs[i].append(image_resized)

            labels = tf.tile(tf.expand_dims(tf.range(n_way),-1), (1, n_query))
            out_imgs = tf.convert_to_tensor(out_imgs)

            return out_imgs, labels


        dataset = dataset.map(
                lambda e_class: tuple(tf.py_func(_parse_py_function, [e_class], [tf.string])))

        dataset = dataset.map(precess_imgs)

        iterator = dataset.make_one_shot_iterator()
        out_imgs, labels = iterator.get_next()
        support, query = tf.split(out_imgs, [n_shot,n_query], 1)

        return support, query, labels