import tensorflow as tf


def main():
    # train(4, [[.1, .2, .3], [.3, .2, .3], [.3, .3, .3], [.3, .4, .5]], [1,2,3,0], [[.1, .2, .3], [.3, .2, .3], [.3, .3, .3], [.3, .4, .5]], [1,2,3,0])
    pass
    '''
    mnist = tf.keras.datasets.mnist

    (x_train, y_train), (x_test, y_test) = mnist.load_data()

    x_train = tf.keras.utils.normalize(x_train, axis=1)
    x_test = tf.keras.utils.normalize(x_test, axis=1)

    print(x_train[0])

    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.Flatten())
    model.add(tf.keras.layers.Dense(128, activation=tf.nn.relu))
    model.add(tf.keras.layers.Dense(128, activation=tf.nn.relu))
    model.add(tf.keras.layers.Dense(10, activation=tf.nn.softmax))

    model.compile(optimizer='SGD',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    model.fit(x_train, y_train, epochs=3)
    '''


# data and results arrays (training and testing) should be paired. Classifications is number of ways to classify data.
def train(classifications: int, data: list, results: list, testdata: list, testresults: list):
    numberOfNeurons = (len(data[0]) + classifications)/2
    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.Flatten())
    model.add(tf.keras.layers.Dense(numberOfNeurons, activation=tf.nn.relu))
    model.add(tf.keras.layers.Dense(numberOfNeurons, activation=tf.nn.relu))
    model.add(tf.keras.layers.Dense(classifications, tf.nn.softmax))

    model.compile(optimizer='SGD',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    model.fit(data, results, epochs=5)

    loss, accuracy = model.evaluate(testdata, testresults)
    print(loss)
    print(accuracy)


if __name__ == '__main__':
    main()