import pickle
import os
import h5py
import numpy as np
from keras.models import load_model, Model
from keras.layers import Input, Flatten, Dense, Dropout
from keras.layers.normalization import BatchNormalization
from keras.applications.vgg16 import VGG16


def initialize_model():

        # Custom input layer
        input = Input(shape=(1036800,), name='image_input')

        # Load convolutional block layers of the VGG16 model
        initial_model = VGG16(weights='imagenet', include_top=False)

        # Add top layers to combine features and predict continuous values
        x = Flatten()(initial_model(input).output)
        x = Dense(200, activation='relu')(x)
        x = BatchNormalization()(x)
        x = Dropout(0.5)(x)
        x = Dense(1)(x)

        # Make new model and compile it
        model = Model(inputs=input, outputs=x)
        model.compile(loss='mse', optimizer='adam')

        print(model.summary())

        return model


def train_model():
        model.fit(x_train, y_train, epochs=20, batch_size=16)
        score = model.evaluate(x_test, y_test, batch_size=16)
        print "Evaluation score: {}".format(score)
        model.save('model.h5')


def fetch_data():
        with h5py.File('/home/aicg2/data/data.h5', 'r') as f:
                x_train, x_test = f['aic480']['train'][:], f['aic480']['val'][:]

        y_train = np.genfromtxt('/home/aicg2/group2/scripts/aic480_train_labels.txt', dtype=int)
        y_test= np.genfromtxt('/home/aicg2/group2/scripts/aic480_val_labels.txt', dtype=int)

        return x_train, x_test, y_train, y_test


if __name__ == '__main__':
        #x_train, x_test, y_train, y_test = fetch_data()
        x_train, x_test, y_train, y_test = pickle.load(open('sample_data.bin', 'rb'))

        if not os.path.isfile('model.h5'):
                model = initialize_model()
        else:
                model = load_model('model.h5')

        train_model()