from chunks1 import create_chunklist
import numpy
numpy.random.seed(123)
import cPickle as pickle
import sys
import io



from keras.engine import Input
from keras.layers.embeddings import Embedding
from keras.layers.recurrent import LSTM
from keras.models import Sequential, Model
from keras.layers import merge
from keras.layers.core import Dense, Dropout, Activation, Reshape, Flatten, Dropout, Highway
from keras.layers.convolutional import Convolution2D, MaxPooling2D, Convolution1D, MaxPooling1D, AveragePooling2D, ZeroPadding2D
from keras.optimizers import SGD, RMSprop, Adagrad, Adadelta, Adam, Adamax, K, Nadam, Optimizer


def train_on_chunklist(model, chunklist):
    result = 0
    for files in chunklist:
        values = []
        labels = []
        for file in files:
            data = pickle.load(open(file, 'rb'))
            values.append(data[0])
            labels.append(data[1])

        chunk_values = numpy.concatenate(values)
        chunk_labels = numpy.concatenate(labels)
        acc = model.train_on_batch(chunk_values, chunk_labels)
        result += acc[1]

    return result / len(chunklist)

def test_on_chunklist(model, chunklist):
    result = 0
    for files in chunklist:
        values = []
        labels = []
        for file in files:
            data = pickle.load(open(file, 'rb'))
            values.append(data[0])
            labels.append(data[1])

        chunk_values = numpy.concatenate(values)
        chunk_labels = numpy.concatenate(labels)
        acc = model.test_on_batch(chunk_values, chunk_labels)
        result += acc[1]

    return result / len(chunklist)


numEpochs = 30

last_train=0
this_train=0
train_count=0

train_limit=20


if __name__ == "__main__":
    inputs = Input(shape=(512, 86), name="Input")
    print(inputs._keras_shape)
    inputs_reshaped = Reshape((1, 512, 86), name="Input reshape")(inputs)
    output = inputs_reshaped
    print(output._keras_shape)
    output = Convolution2D(512, 15, 86, input_shape='1,512,86', border_mode='valid', activation='relu', name="Convolution 512x15x86")(output)
    print(output._keras_shape)
    output = MaxPooling2D(pool_size=(2, 1), name="MAxPool1 2x1")(output)
    print(output._keras_shape)
    print(inputs_reshaped._keras_shape)
    input_img = output
    tower_1 = Convolution2D(512, 1, 1, border_mode='same', activation='relu', name="Convolution1 512x1x1")(input_img)
    print(tower_1._keras_shape)
    tower_1 = Convolution2D(512, 3, 1, border_mode='same', activation='relu', name="Convolution1 512x3x1")(tower_1)
    print(tower_1._keras_shape)
    print("1")
    tower_2 = Convolution2D(512, 1, 1, border_mode='same', activation='relu', name="Convolution2 512x1x1")(input_img)
    print(tower_2._keras_shape)
    tower_2 = Convolution2D(512, 5, 1, border_mode='same', activation='relu', name="Convolution2 512x5x1")(tower_2)
    print(tower_2._keras_shape)
    print("2")
    tower_3 = MaxPooling2D((3, 1), strides=(1, 1), border_mode='same', name="MAxPool1 3,3")(input_img)
    print(tower_3._keras_shape)
    tower_3 = Convolution2D(512, 15, 1, border_mode='same', activation='relu', name="Convolution3 512x15x1")(tower_3)
    print(tower_3._keras_shape)
    print("3")
    output = merge([tower_1, tower_2, tower_3, input_img], mode='concat', concat_axis=1, name="Merge Inception 1")
    print(output._keras_shape)
    output = MaxPooling2D(pool_size=(2, 1), name="MaxPool2 2x1")(output)
    output = Convolution2D(1024, 15, 1, border_mode='valid', activation='relu', name="Convolution1 1024x15x1")(output)
    print(output._keras_shape)
    output = MaxPooling2D(pool_size=(2, 1), name="MaxPool3 2x1")(output)
    print(output._keras_shape)
    input_img = output
    tower_1 = Convolution2D(512, 1, 1, border_mode='same', activation='relu', name="Convolution4 512x1x1")(input_img)
    print(tower_1._keras_shape)
    tower_1 = Convolution2D(512, 3, 1, border_mode='same', activation='relu', name="Convolution4 512x3x1")(tower_1)
    print(tower_1._keras_shape)
    print("1")
    tower_2 = Convolution2D(512, 1, 1, border_mode='same', activation='relu', name="Convolution5 512x1x1")(input_img)
    print(tower_2._keras_shape)
    tower_2 = Convolution2D(512, 5, 1, border_mode='same', activation='relu', name="Convolution5 512x5x1")(tower_2)
    print(tower_2._keras_shape)
    print("2")
    tower_3 = MaxPooling2D((3, 1), strides=(1, 1), border_mode='same', name="MAxPool2 3,3")(input_img)
    print(tower_3._keras_shape)
    tower_3 = Convolution2D(512, 15, 1, border_mode='same', activation='relu', name="Convolution6 512x15x1")(tower_3)
    print(tower_3._keras_shape)
    print("3")
    output = merge([tower_1, tower_2, tower_3, input_img], mode='concat', concat_axis=1, name="Merge Inception 2")
    print(output._keras_shape)
    output = MaxPooling2D(pool_size=(2, 1))(output)
    output = Convolution2D(1024, 15, 1, border_mode='valid', activation='relu', name="Convolution2 1024x15x1")(output)
    print(output._keras_shape)
    output = MaxPooling2D(pool_size=(2, 1), name="MaxPool4 2x1")(output)
    print(output._keras_shape)
    #output = Convolution2D(350, 5, 1)(output)
    #print(output._keras_shape)
    #output = MaxPooling2D(pool_size=(2, 1))(output)
    #print(output._keras_shape)
    output = Flatten(name="Flatten")(output)
    print(output._keras_shape)
    output = Dense(2048, activation='relu', name="Dense 2048")(output)
    print(output._keras_shape)
    output = Dropout(0.125, name="Droupout1 0.125")(output)
    output = Dense(1024, activation='relu', name="Dense 1024")(output)
    print(output._keras_shape)
    output = Dropout(0.125, name="Droupout2 0.125")(output)
    output = Dense(2, activation='softmax', name="Dense Softmax output")(output)
    print(output._keras_shape)
        


    model = Model(input=inputs, output=output)

    sgd = SGD(lr=0.001, decay=1e-8, momentum=0.9, nesterov=True)
    model.compile(loss='binary_crossentropy', optimizer=sgd, metrics=['accuracy'])

    train_chunklist = create_chunklist('train')
    test_chunklist = create_chunklist('test')

    open("NN23_log.txt", "wt")

    for epoch in range(1, numEpochs+1):
        epoch_result = train_on_chunklist(model, train_chunklist)
	this_train=epoch_result
	if '%.4f' % this_train == '%.4f' % last_train:
		train_count +=1
	else:
		train_count=0
	last_train = this_train
	if train_count>train_limit:
		sys.exit(0)
        #print "EPOCH {0} result = {1}".format(epoch, '%.6f' % epoch_result)
	open("NN23_log.txt", "a").write("EPOCH {0} result {1} ".format(epoch, '%.6f' % epoch_result))
	epoch_test = test_on_chunklist(model, test_chunklist)
    	#print "EPOCH {0} test = {1}".format(epoch, '%.6f' % epoch_test)
	open("NN23_log.txt", "a").write("test {0} \n".format('%.6f' % epoch_test))
