"""train_classifier.py
Train_classifier will load the training and testing
data, build the neural network graph, then train. It will
return and save the model as a .hdf5 file
"""
import pathlib
import argparse
import json
from keras.models import load_model
from keras.callbacks import History, ModelCheckpoint, EarlyStopping
from keras.preprocessing.image import ImageDataGenerator
import model

HISTORY = History()
BATCH_SIZE = 256
EPOCHS = 10


def create_feed_data(directory, height, width):
    """
    Create batchs of feeding data by going directly
    to the images. Labels are generated by the file
    structure.
    INPUT: Directory of images. The directory should
    be structured as below:
    Parent
    |- Training
    |   |- Label
    |       |- Image
    |- Testing
        |- Label
            |- Image
    OUTPUT: Keras training and validation generators
    """
    # Create the directory variables
    dataset_path = pathlib.Path(directory)
    train_path = dataset_path.joinpath("Training")
    test_path = dataset_path.joinpath("Testing")

    img_width, img_height = width, height

    # Obtain the number of images in the both
    # training and testing directory
    no_train_sam = 0
    no_test_sam = 0

    for label in train_path.iterdir():
        for image in label.iterdir():
            no_train_sam += 1

    for label in test_path.iterdir():
        for image in label.iterdir():
            no_test_sam += 1

    # Data augmentation (Rescale the image values)
    train_datagen = ImageDataGenerator(rescale=1. / 255,
                                       featurewise_center=True,
                                       featurewise_std_normalization=True,
                                       rotation_range=25,
                                       shear_range=0.3,
                                       zoom_range=0.3,
                                       width_shift_range=0.1,
                                       height_shift_range=0.1,
                                       fill_mode="constant",
                                       cval=0,
                                       zca_whitening=True)
    test_datagen = ImageDataGenerator(rescale=1. / 255)

    # Flow the images directly from the directory
    train_gen = train_datagen.flow_from_directory(
        train_path,
        target_size=(img_width, img_height),
        color_mode="grayscale",
        batch_size=BATCH_SIZE,
        class_mode='categorical')
    test_gen = test_datagen.flow_from_directory(
        test_path,
        target_size=(img_width, img_height),
        color_mode="grayscale",
        batch_size=BATCH_SIZE,
        class_mode='categorical')

    # Create the dictionary and write to a file for the classifier
    lab_dict = (test_gen.class_indices)
    # Flip the key-values
    lab_dict = dict((v, k) for k, v in lab_dict.items())

    return train_gen, test_gen, no_train_sam, no_test_sam, lab_dict


def build_model(classes, height, width):
    """Build the Keras neural network"""
    print("> Building Keras neural network...")
    network_model = model.simple_3(classes=classes, height=height, width=width)
    return network_model


def fit_model(model, train_gen, test_gen, no_train, no_test, name):
    """
    DESCRIPTION: Train the neural network with batch size of BATCH, 10 epochs,
    and print out training progress. Uses checkpoints and early stopping.
    INPUT: Keras Training and Testing generators, number of samples, and model
    OUTPUT: Keras Trainined model
    """
    # Checkpoints for each iteration.
    filename = 'weights.{epoch:02d}-{val_loss:.2f}-{val_acc:.2f}.hdf5'
    checkpointer = ModelCheckpoint(filepath=filename, verbose=1, period=1)

    # Early stopping to stop training when accuracy is not improving.
    early_stopping = EarlyStopping(monitor='val_loss',
                                   min_delta=0,
                                   patience=2,
                                   verbose=0,
                                   mode='auto')

    # Begin training with the two callbacks.
    model.fit_generator(
        train_gen,
        steps_per_epoch=no_train // BATCH_SIZE,
        epochs=EPOCHS,
        validation_data=test_gen,
        validation_steps=no_test // BATCH_SIZE,
        callbacks=[checkpointer, early_stopping]
    )

    model.save(name)
    return model


if __name__ == "__main__":
    # Build up the argument to bring in an image
    AP = argparse.ArgumentParser()
    AP.add_argument("-m", "--model_path", help="path to an existing model",
                    required=False)
    AP.add_argument("-d", "--data_path", help="path to dataset", required=True)
    AP.add_argument("-s", "--save_path", help="path to save the final model",
                    required=True)
    AP.add_argument("-c", "--classes", help="number of classes", required=True)
    ARGS = vars(AP.parse_args())

    #############################
    # TRAIN
    #############################
    DIRECTORY = ARGS['data_path']
    MODEL_PATH = ARGS['model_path']
    MODEL_SAVE_NAME = ARGS['save_path']
    CLASSES = int(ARGS['classes'])
    HEIGHT = 32
    WIDTH = 32
    DICT_PATH = '/home/alexander/Desktop/projects/ErmisOCR/src/classifier/char_labels' + str(CLASSES) + '.json'
    # Create training and testing data generators
    TRAIN_GEN, TEST_GEN, NO_TRAIN, NO_TEST, DICT = create_feed_data(DIRECTORY,
                                                                    HEIGHT,
                                                                    WIDTH)

    # Print out the label dictionary
    with open(DICT_PATH, 'w') as file:
        json.dump(DICT, file)

    # Check if the model already exists
    # If not, build a new model file
    # If it does exist, load the model
    if MODEL_PATH is None:
        CHAR_MODEL = build_model(CLASSES, HEIGHT, WIDTH)
    else:
        MODEL_DIRECTORY = pathlib.Path(MODEL_PATH)
        CHAR_MODEL = load_model(MODEL_DIRECTORY)

    HISTORY = fit_model(CHAR_MODEL,
                        TRAIN_GEN,
                        TEST_GEN,
                        NO_TRAIN,
                        NO_TEST,
                        MODEL_SAVE_NAME)
