import numpy as np
import sys, os, re, gzip, struct
import random
import h5py
import copy
from keras import backend as K
from keras.utils import Sequence
import keras.utils
import tensorflow as tf

ASC_CLASS=10

class ASCDataGenerator(Sequence):

    def __init__(self, file, keys=None, batch_size=64, dim=(40,500), n_channels=1, shuffle=True, mode=None,
                 mean=None, var=None):

        self.file=file
        self.batch_size=batch_size
        self.dim=dim
        self.n_classes=ASC_CLASS
        self.n_channels=n_channels
        self.shuffle=shuffle
        self.mode=mode
        self.mean=mean
        self.var=var
        self.keys=[]

        if keys is not None:
            self.keys=keys

        self.true_labels={"airport":0, "bus":1, "metro":2, "metro_station":3, "park":4, "public_square":5,
                          "shopping_mall":6, "street_pedestrian":7, "street_traffic":8, "tram":9}

        self.h5fd = h5py.File(self.file, 'r')
        if keys is None:
            self.n_samples = len(self.h5fd.keys())
            for key in self.h5fd.keys():
                self.keys.append(key)
        else:
            self.n_samples = len(self.keys)

        if self.shuffle:
            random.shuffle(self.keys)

    def __len__(self):
        return int(np.ceil(self.n_samples)/self.batch_size)
        #return int(np.floor(self.n_samples / self.batch_size))


    def __getitem__(self, index):
        list_keys_temp = [self.keys[k] for k in range(index*self.batch_size,
                                                      min( (index+1)*self.batch_size, len(self.keys) ) )]

        X, y = self.__data_generation(list_keys_temp)

        if self.mode == 'double' :
            return X, [y,y]
        elif self.mode == 'reverse':
            y_rev=np.flip(y, 0)
            y_rev=np.max(np.multiply(y, y_rev), axis=1, keepdims=False)
            y_rev=np.reshape(y_rev, (self.batch_size, 1))
            return X, [y, y_rev]
        else:
            return X, y

    def on_epoch_end(self):
        if self.shuffle == True:
            random.shuffle(self.keys)

    def get_label(self, key):
        #spk=key.split('-')
        return self.true_labels[key.split('-')[0]]

    def __data_generation(self, list_keys_temp):
        x = np.zeros((self.batch_size, *self.dim, self.n_channels))
        y = np.empty((self.batch_size), dtype=int)

        self.mean=np.reshape(self.mean, (self.dim[0],1))
        self.var=np.reshape(self.var, (self.dim[0],1))

        for i, key in enumerate(list_keys_temp):
            mat = self.h5fd[key][()]
            # normalization
            mat = mat - self.mean
            mat = np.divide(mat, self.var)

            if mat.shape[1] > self.dim[1]:
                mat=mat[:, 0:self.dim[1]]

            #src = mat.transpose([1,0])
            src = np.reshape(mat, (1, self.dim[0], self.dim[1], 1))
            x[i,:,0:src.shape[2],:] = src
            y[i] = self.get_label(key)

        return x, keras.utils.to_categorical(y, num_classes=self.n_classes)

    def compute_norm(self):
        mean=None
        sq_mean=None
        rows=0
        for key in self.keys:
            mat=self.h5fd[key][()]
            rows += mat.shape[1]
            if mean is None:
                mean=np.sum(mat, axis=1).astype(np.float64)
                sq_mean=np.sum(np.square(mat), axis=1).astype(np.float64)
            else:
                mean=np.add(np.sum(mat, axis=1).astype(np.float64), mean)
                sq_mean=np.add(np.sum(np.square(mat), axis=1).astype(np.float64), sq_mean)
        mean = mean/rows
        sq_mean = sq_mean/rows - np.square(mean)
        sq_mean=np.sqrt(sq_mean+1.0e-8)

        self.mean=mean.astype(np.float32)
        self.var=sq_mean.astype(np.float32)

    def get_mean_var(self):
        return self.mean, self.var

    def save_norm(self, file):
        with h5py.File(file, 'w') as f:
            f.create_dataset('mean', data=self.mean, compression='gzip', compression_opts=9)
            f.create_dataset('var', data=self.var, compression='gzip', compression_opts=9)

    def load_norm(self, file):
        with h5py.File(file, 'r') as f:
            self.mean=f['mean'][()]
            self.var=f['var'][()]
'''
class ASCBatchGenerator(object):
    def __init__(self, file, keys=None, batch_size=64, dim=(40,500), n_channels=1, shuffle=True, mode=False,
                 mean=None, var=None):
        self.file=file
        self.batch_size=batch_size
        self.n_classes=ASC_CLASS
        self.dim=dim
        self.src_list=[]
        self.dst_list=[]
        self.mean=mean
        self.var=var
        self.keys=[]

        if keys is not None:
            self.keys = keys

        self.h5fd=h5py.File(self.file, 'r')
        if keys is None:
            self.n_samples = len(h5fd.keys())
        else:
            self.n_samples = len(self.keys)

        if keys is None:
            for key in self.h5fd.keys():
                self.keys.append(key)

        self.shuffle=shuffle
        if self.shuffle:
            random.shuffle(self.keys)

        self.mode=mode

        self.sample_per_epoch = int(np.ceil(self.n_samples / self.batch_size))

        self.true_labels={"airport":0, "bus":1, "metro":2, "metro_station":3, "park":4, "public_square":5,
                          "shopping_mall":6, "street_pedestrian":7, "street_traffic":8, "tram":9}

    def reset(self):
        self.src_list = []
        self.dst_list = []

    def get_label(self, key):
        spk=key.split('-')
        return self.true_labels[spk[0]]

    def get_sample_per_epoch(self):
        return self.sample_per_epoch

    def get_num_samples(self):
        return self.n_samples

    def gen_batch(self):

        if self.shuffle:
            random.shuffle(self.keys)

        for key in self.keys:
            mat=self.h5fd[key][()]
            # normalization
            mat = mat-self.mean
            mat = np.divide(mat, self.var)

            src=mat.transpose([1,0])
            src=src.reshape(1, src.shape[0], src.shape[1], 1)

            label=self.get_label(key)

            self.src_list.append(src)
            self.dst_list.append(label)

            if len(self.src_list) == self.batch_size:
                inputs = np.zeros((self.batch_size, *self.dim, 1))
                outputs = np.empty((self.batch_size), dtype=int)

                for n in range(self.batch_size):
                    inputs[n,:,0:self.src_list[n].shape[2],:] = self.src_list[n]
                    outputs[n] = self.dst_list[n]
                self.reset()

                label=keras.utils.to_categorical(outputs, num_classes=self.n_classes)
                if self.mode:
                    yield inputs, [label, label]
                else:
                    yield inputs, label

        # remained objects
        if len(self.src_list) > 0:
            inputs = np.zeros((self.batch_size, *self.dim, 1))
            outputs = np.empty((self.batch_size), dtype=int)

            for n in range(len(self.src_list)):
                inputs[n,:,0:self.src_list[n].shape[2],:] = self.src_list[n]
                outputs[n] = self.dst_list[n]
                self.reset()

            label=keras.utils.to_categorical(outputs, num_classes=self.n_classes)
            if self.mode:
                yield inputs, [label, label]
            else:
                yield inputs, label

    def compute_norm(self):
        mean=None
        sq_mean=None
        rows=0
        for key in self.keys:
            mat=self.h5fd[key][()]
            rows+=mat.shape[0]
            if mean is None:
                mean=np.sum(mat, axis=0).astype(np.float64)
                sq_mean=np.sum(np.square(mat), axis=0).astype(np.float64)
            else:
                mean=np.add(np.sum(mat, axis=0).astype(np.float64), mean)
                sq_mean=np.add(np.sum(np.square(mat), axis=0).astype(np.float64), sq_mean)
        mean = mean/rows
        sq_mean = sq_mean/rows - np.square(mean)
        sq_mean=np.sqrt(sq_mean+1.0e-8)

        self.mean=mean.astype(np.float32)
        self.var=sq_mean.astype(np.float32)

    def get_mean_var(self):
        return self.mean, self.var

    def save_norm(self, file):
        with h5py.File(file, 'w') as f:
            f.create_dataset('mean', data=self.mean, compression='gzip', compression_opts=9)
            f.create_dataset('var', data=self.mean, compression='gzip', compression_opts=9)

    def load_norm(self, file):
        with h5py.File(file, 'r') as f:
            self.mean=f['mean'][()]
            self.var=f['var'][()]

'''
