import tensorflow as tf
import numpy as np
import data_helper

class TLU(object):
    """
        Thresholding Logic Unit class definition
    """
    graph = None
    weight = np.asarray(
        [[1.0],
         [-1.0]]
    )

    def __init__(self):
        self.graph = tf.Graph()
        with self.graph.as_default():
            # Define input and weight
            self.weight = tf.Variable(self.weight, dtype=tf.float16)
            self._input = tf.placeholder(tf.float16, [None, 2], name='input')

            # multiplication
            mul = tf.matmul(self._input, self.weight)

            # Step function
            self._output = tf.where( tf.less(mul, tf.zeros(tf.shape(mul), dtype=tf.float16)),
                tf.zeros(tf.shape(mul)),
                tf.ones(tf.shape(mul)))

    def predict(self, x):
        """
            Predict the result
        """
        with tf.Session(graph=self.graph) as sess:
            sess.run(tf.global_variables_initializer())
            return sess.run([self._output,], feed_dict={self._input: x})

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

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