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 = 20

last_train=0
this_train=0
train_count=0

train_limit=20


if __name__ == "__main__":
    inputs = Input(shape=(512, 86))
    print(inputs._keras_shape)
    inputs_reshaped = Reshape((1, 512, 86))(inputs)
    output = inputs_reshaped
    print(output._keras_shape)
    output = Convolution2D(100, 7, 86, input_shape='1,512,86', border_mode='valid', activation='relu')(output)
    print(output._keras_shape)
    output = MaxPooling2D(pool_size=(2, 1))(output)
    print(output._keras_shape)
    output = Convolution2D(150, 4, 1, border_mode='valid', activation='relu')(output)
    print(output._keras_shape)
    output = MaxPooling2D(pool_size=(2, 1))(output)
    print(output._keras_shape)
    output = Convolution2D(250, 4, 1, border_mode='valid', activation='relu')(output)
    print(output._keras_shape)
    output = MaxPooling2D(pool_size=(2, 1))(output)
    print(output._keras_shape)
    output = Flatten()(output)
    print(output._keras_shape)
    output = Dense(512, activation='relu')(output)
    print(output._keras_shape)
    output = Dropout(0.125)(output)
    output = Dense(512, activation='relu')(output)
    print(output._keras_shape)
    output = Dropout(0.125)(output)
    output = Dense(2, activation='softmax')(output)
    print(output._keras_shape)

    

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

    #sgd = SGD(lr=0.001, decay=1e-8, momentum=0.9, nesterov=True)
    adadelta = Adadelta(lr=1.0, rho=0.95, epsilon=1e-08, decay=0.0)
    model.compile(loss='binary_crossentropy', optimizer=adadelta, metrics=['accuracy'])


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

    open("NN31_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("NN31_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("NN31_log.txt", "a").write("test {0} \n".format('%.6f' % epoch_test))
