from keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img
import numpy as np
import pandas as pd

import keras

print('done importing')

def completeWithAugmentedData(X_tr,Y_tr,limit):
    print('complete augemented data')
    toAdd=limit-len(X_tr)
    if (len(X_tr)<limit):
        datagen = ImageDataGenerator(
                shear_range=0.1,
                zoom_range=0.1,
                horizontal_flip=True)
    gentrain = datagen.flow(np.asarray(X_tr), batch_size=1)
    igen=0
    for gt in gentrain:
        X_tr.append(gt[0,:,:,:])
        Y_tr.append(Y_tr[0])
        igen += 1
        if (igen >= toAdd):
            break
        #print len(X_tr),len(Y_tr)
    return [X_tr,Y_tr]

def read_data_small(labels_dim):

    print('run read_data_small')
    input_dir="/mnt/server-home/TUE/20184102/datasets/train"
    labels=pd.read_csv("/mnt/server-home/TUE/20184102/datasets/trainLabels.csv")

    X_train0 = []
    X_train1 = []

    Y_train0 = []
    Y_train1 = []

    limit=100

    for idx,filename in enumerate(labels["image"]):
        img = load_img(input_dir+"/"+ filename+".jpeg", target_size=(256, 256)) # this is a PIL image
        x = img_to_array(img) # this is a Numpy array with shape (128, 128, 3)
        print('start for loop')
        if ((labels["level"][idx]==0)and(len(X_train0)<limit)):
            X_train0.append(x)
            Y_train0.append(labels["level"][idx])
            print('3')

        if (((labels["level"][idx]==1)) or (labels["level"][idx]==2) or (labels["level"][idx]==3) or (labels["level"][idx]==4)) and (len(X_train1)<limit):
            X_train1.append(x)
            Y_train1.append(labels["level"][idx])
            print('2')

        if (len(X_train1)>= limit) and (len(X_train0) >= limit):
            break


    [X_train0,Y_train0]=completeWithAugmentedData(X_train0,Y_train0,limit)
    [X_train1, Y_train1] = completeWithAugmentedData(X_train1, Y_train1,limit)

    X_train=np.concatenate((np.asarray(X_train0),np.asarray(X_train1)),axis=0)
    Y_train=np.concatenate((np.asarray(Y_train0),np.asarray(Y_train1)),axis=0)

    X_train = X_train.astype('float16') / 255.
    Y_train = keras.utils.to_categorical(Y_train, labels_dim)

    print ("TrainData size", X_train.shape, Y_train.shape)
    np.savez_compressed("data/trainDataMedium.npz",X_train=X_train,Y_train=Y_train)

read_data_small(2)

#read_data_small_test(5)

"""
def read_data(labels_dim):
  input_dir="data/trainData"
  labels=pd.read_csv("data/trainData.csv")

  X_train=[]
  Y_train=[]

  for idx,filename in enumerate(labels["image"]):
    img = load_img(input_dir+"/"+str(labels["level"][idx])+"/"+ filename+".tiff") # this is a PIL image
    x = img_to_array(img) # this is a Numpy array with shape (128, 128, 3)
    X_train.append(x)
    Y_train.append(labels["level"][idx])

  X_train=np.asarray(X_train)
  Y_train=np.asarray(Y_train)
  X_train = X_train.astype('float16') / 255.
  Y_train = keras.utils.to_categorical(Y_train, labels_dim)

  print ("TrainData size", X_train.shape, Y_train.shape)
  np.savez_compressed("data/trainData.npz",X_train=X_train,Y_train=Y_train)




def read_train_data():
  data=np.load("data/trainData.npz")
  X_train=data["X_train"]
  Y_train=data["Y_train"]
  return [X_train,Y_train]

def read_test_data():
  data=np.load("data/testData.npz")
  X_test=data["X_test"]
  Y_test=data["Y_test"]
  return [X_train,Y_train]

[X_train,Y_train]=read_train_data()
print "bau"
for j in range (100000000):
  c=j+j*1.2321
[X_test,Y_test]=read_test_data()
"""