import pandas as pd
import numpy
import tensorflow as tf
from tensorflow.keras import datasets, layers, models
import math
import time
import tkinter
from matplotlib import pyplot as plt


df = pd.read_csv('train.csv')
df = df.sample(frac=1)

df_label = df['label']
del df['label']

X = df.to_numpy()
y = df_label.to_numpy()


X = X.reshape((-1, 28, 28, 1))
X_train_size = math.floor(len(X)*.8)
X_test_size = len(X) - X_train_size

X_train = X[:X_train_size]
X_test = X[X_test_size+1:]

y_train = y[:X_train_size]
y_test = y[X_test_size+1:]


y_train = numpy.eye(10)[y_train]
y_test = numpy.eye(10)[y_test]

X_train = X_train/255
X_test = X_test/255

print(X_train[0])


model = models.Sequential()
model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(128, (3, 3), activation='relu'))



model.add(layers.Flatten())
model.add(layers.Dense(64, activation='relu'))
model.add(layers.Dense(10, activation='softmax'))


optmizer = tf.keras.optimizers.Adam(learning_rate=0.001)
model.compile(optimizer=optmizer, loss=tf.keras.losses.MeanSquaredError(),
            metrics=['accuracy'])
history = model.fit(X_train, y_train, batch_size=2**12, epochs=30, validation_data=(X_test, y_test))



pred = model.predict(X_test)

#print(numpy.argmax(pred[255]))
#plt.imshow(X_test[255].reshape(28,28))
#plt.show()

df = pd.read_csv('test.csv')
X = df.to_numpy()

y = model.predict(X.reshape(-1,28,28,1))

with open('submission.csv', 'w') as f:
    f.write("ImageId,Label\n")
    idx=1
    for i in y:
        f.write(f"{idx},{numpy.argmax(i)}\n")
        idx+=1