import numpy as np
import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Conv2D, AvgPool2D, Flatten, Dense, Dropout


def load_train(path):
    x_train = np.load(path + 'train_features.npy')
    y_train = np.load(path + 'train_target.npy')
    x_train = x_train.reshape(x_train.shape[0],28,28,1)/255.0
    return x_train, y_train 
    
def create_model(input_shape):
    model = Sequential()
    model.add(Conv2D(filters = 32, kernel_size=(3,3),padding = 'same', activation = 'relu',input_shape=(28,28,1)))
    model.add(AvgPool2D(pool_size=(2, 2)))
    model.add(Conv2D(filters = 64, kernel_size=(3,3), activation = 'relu'))
    model.add(AvgPool2D(pool_size=(2, 2)))
    model.add(Conv2D(filters = 128, kernel_size=(3,3), activation = 'relu'))
    model.add(Flatten())
    model.add(Dense(units=100,activation = 'relu'))
    model.add(Dense(units=50,activation = 'relu'))
    model.add(Dense(units=30,activation = 'relu'))
    model.add(Dense(units=10,activation = 'softmax'))
    model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['acc'])
    return model
    
def train_model(model, train_data, test_data, batch_size=32, epochs=5,
                steps_per_epoch=None, validation_steps=None): 
    x_train, y_train = train_data
    x_test, y_test = test_data
    model.fit(train_data, validation_data=test_data,batch_size=batch_size, 
              epochs=epochs, steps_per_epoch=steps_per_epoch,validation_steps=validation_steps,
              verbose=2, shuffle=True)
    return model