import numpy as np
import matplotlib.pyplot as plt
import cv2 as cv
import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras import optimizers
from tensorflow.keras import utils
from tqdm import tqdm

# load data
mnist = tf.keras.datasets.mnist
(x_train, y_train),(x_test, y_test) = mnist.load_data()

# prepare data
x_train, x_test = x_train / 255.0, x_test / 255.0 # normalize
x_train = x_train.reshape(-1, 28, 28)
x_test = x_test.reshape(-1, 28, 28)
y_train_oh = utils.to_categorical(y_train, num_classes = 10) # one-hot encoding
y_test_oh = utils.to_categorical(y_test, num_classes = 10) # one-hot encoding

TIME_STEP = 28 # same sa the height of image
INPUT_SIZE = 28 # same as the width of image
BATCH_SIZE = 256
BATCH_INDEX = 0
OUTPUT_SIZE = 10 # one-hot
CELL_SIZE = 50 # number of hidden units
LR = 0.001

# build RNN model
model = tf.keras.models.Sequential([
    layers.SimpleRNN(CELL_SIZE, activation='tanh'),
    # output layer
    layers.Dense(OUTPUT_SIZE, activation='softmax')
])

# specified optimizer
adam = optimizers.Adam(LR)
# compile model
model.compile(optimizer=adam, 
             loss='categorical_crossentropy',
             metrics=['accuracy'])

# start to train model
for step in range(10001):
    x_batch = x_train[BATCH_INDEX: BATCH_INDEX+BATCH_SIZE, :, :]
    y_batch = y_train_oh[BATCH_INDEX: BATCH_INDEX+BATCH_SIZE, :]
    cost = model.train_on_batch(x_batch, y_batch)
    
    BATCH_INDEX += BATCH_SIZE
    BATCH_INDEX = 0 if BATCH_INDEX >= x_train.shape[0] else BATCH_INDEX
    
    if step % 1000 == 0:
        cost, acc = model.evaluate(x_test, y_test_oh, batch_size=y_test.shape[0], verbose=False)
        print('test cost: ', cost, 'test accuracy: ', acc)

test cost:  2.275296926498413 test accuracy:  0.1778
test cost:  0.3357826769351959 test accuracy:  0.9023
test cost:  0.2307487279176712 test accuracy:  0.9347
test cost:  0.18899917602539062 test accuracy:  0.9472
test cost:  0.16296234726905823 test accuracy:  0.9533
test cost:  0.16683271527290344 test accuracy:  0.9528
test cost:  0.14237910509109497 test accuracy:  0.9608
test cost:  0.15203972160816193 test accuracy:  0.96
test cost:  0.14567141234874725 test accuracy:  0.9592
test cost:  0.13654865324497223 test accuracy:  0.9628
test cost:  0.1286523938179016 test accuracy:  0.9651

model.fit(x_train, y_train_oh, epochs=5, batch_size=None) # If unspecified, batch_size will default to 32.
model.evaluate(x_test, y_test_oh)