from tensorflow.keras.datasets import fashion_mnist
from tensorflow.keras.layers import Conv2D, Flatten, Dense, MaxPooling2D
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import numpy as np



def load_train(path):
    features_train = np.load(path + 'train_features.npy')
    target_train = np.load(path + 'train_target.npy')
    features_train = features_train / 255
    features_train = np.expand_dims(features_train, axis=3)

    return features_train, target_train

def load_test(path):
    features_test = np.load(path + 'test_features.npy')
    target_test = np.load(path + 'test_target.npy')
    features_test = features_test / 255
    features_test = np.expand_dims(features_test, axis=3)

    return features_test, target_test


def create_model(input_shape):
    model = Sequential()
    model.add(Conv2D(filters=4, kernel_size=(3, 3), padding='same',
                 activation="relu", input_shape=input_shape))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Flatten())
    model.add(Dense(units=10, activation='softmax'))
    optimizer = Adam(lr=0.005)
    model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy',
                  metrics=['acc'])


    return model



def train_model(model, train_data, test_data, batch_size=None, epochs=47,
               steps_per_epoch=None, validation_steps=None):


    features_train, target_train = train_data
    features_test, target_test = test_data
    model.fit(features_train, target_train,
              validation_data=(features_test, target_test),
              batch_size=batch_size, epochs=epochs,
              steps_per_epoch=steps_per_epoch,
              validation_steps=validation_steps,
              verbose=2, shuffle=True)


    return model