# importing models 
import tensorflow as tf 
import tensorflow_datasets as tfds
from tensorflow.keras.layers import Flatten, Dense
from tensorflow.keras import Sequential

#loading data (as supervised and with data information)
mnist_data, mnist_info = tfds.load(name="mnist", with_info=True, as_supervised=True)

# extracting train and test data 
mnist_train, mnist_test = mnist_data["train"], mnist_data["test"]

#setting validation size (cast as integer)
validation_samples_num = 0.1*mnist_info.splits["train"].num_examples
validation_samples_num = tf.cast(validation_samples_num, tf.int64)

# get test size as well 
test_samples_num = mnist_info.splits["test"].num_examples
test_samples_num = tf.cast(test_samples_num, tf.int64)

#scale data to make numerically stable (0-white to 255-black), scale within 0-1
def scale(image, label):
    #cast the image into float 
    image = tf.cast(image, tf.float32)
    image /= 255.
    return image, label

# use .map method for customized transformeers
train_scaled_data = mnist_train.map(scale)
test_scaled_data = mnist_test.map(scale)

# to prevent the repetitive patterns, shuffling can challenge our model 
BUFFER_SIZE=10000
train_scaled_shuffle = train_scaled_data.shuffle(BUFFER_SIZE)

#extract the 10% for validation using take and skip method 
validation_scaled_shuffle = train_scaled_shuffle.take(validation_samples_num)
trained_scaled_shuffle = train_scaled_shuffle.skip(validation_samples_num)

# set batch size of model (which help to train the model)--helpful for my poor laptop 
BATCH_SIZE = 100
train_data = train_scaled_shuffle.batch(BATCH_SIZE)

#since we only forward propagate to calculate loss in the validation and test data, we set one batch (size of sample)
validation_data = validation_scaled_shuffle.batch(validation_samples_num)
test_data = test_scaled_data.batch(test_samples_num)

#split input and target data 
validation_inputs, validation_targets = next(iter(validation_data))

def build_nn(hidden_layer_size, input_size, output_size, plot_network=False, activation="relu"):
    
    #simple structure---complex like 5000 (10 layers) on 100 epochs will give 99% accuracy 
    model=Sequential()
    model.add(Flatten(input_shape=(28,28,1))) #transform into vector data (input layer)
    model.add(Dense(hidden_layer_size, activation=activation)) #1st layer
    model.add(Dense(hidden_layer_size, activation=activation))#2nd layer
    model.add(Dense(output_size, activation="softmax")) #output layer (since categorical data transform outputs into probabilities)
    
    from ann_visualizer.visualize import ann_viz
    if plot_network ==True:
        ann_viz(model, title="Dense Model")
    
    model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])
    return model 
model = build_nn(hidden_layer_size=50, input_size=784, output_size=10, plot_network=False, activation="relu")