from keras.layers import Dense, Input
from keras.models import Model
import numpy as np
import data_helper

class TLU(object):
    """
        Thresholding Logic Unit class definition
    """
    epoch = None
    model = None

    def __init__(self, epoch=5000):
        self.epoch = epoch
        self.input = Input((2,))
        mul = Dense((1), activation='sigmoid')(self.input)
        self.model = Model(inputs=self.input, outputs=mul)
        self.model.compile(loss='mse', optimizer='adam')

    def train(self, x, y):
        """
            Train the TLU model for specific training epoch
        """
        for i in range(self.epoch):
            self.model.fit(x, y)

    def predict(self, x):
        """
            Predict the result
        """
        return self.model.predict(x)

if __name__ == '__main__':
    # Load data and build model
    train_x, train_y, test_x = data_helper.load()
    cell = TLU()

    # Train
    cell.train(train_x, train_y)

    # Predict
    print "<< testing >>"
    result = cell.predict(test_x)
    for i in range(len(test_x)):
        print "test index: ", i, '\tresult: ', result[i]