#  Copyright 2021 Industrial Technology Research Institute
#
#  Copyright 2020 Xilinx Inc.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
#  NOTICE: This file has been modified by Industrial Technology Research Institute for AIdea "FPGA Edge AI – AOI Defect
#  Classification" competition tutorial
#
#  Original Author: Mark Harvey, Xilinx Inc.
#
#  Training script for AIdea "FPGA Edge AI – AOI Defect Classification" competition tutorial.

import os
import shutil
import sys
import argparse

# Silence TensorFlow messages
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

import tensorflow as tf
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import TensorBoard, ModelCheckpoint, LearningRateScheduler

from dataset_utils import input_fn_trn, input_fn_test
from customcnn import customcnn

DIVIDER = '-----------------------------------------'


def train(input_height, input_width, input_chan, tfrec_dir, batchsize, learnrate, epochs, chkpt_dir, tboard):
    def step_decay(epoch):
        """
        Learning rate scheduler used by callback
        Reduces learning rate depending on number of epochs
        """
        lr = learnrate
        if epoch > 200:
            lr /= 100
        elif epoch > 10:
            lr /= 10
        return lr

    '''
    Define the model
    '''
    model = customcnn(input_shape=(input_height, input_width, input_chan), classes=6, filters=[8, 16, 32, 64, 128])

    print('\n' + DIVIDER)
    print(' Model Summary')
    print(DIVIDER)
    print(model.summary())
    print("Model Inputs: {ips}".format(ips=(model.inputs)))
    print("Model Outputs: {ops}".format(ops=(model.outputs)))

    '''
    tf.data pipelines
    '''
    # train and test folders
    train_dataset = input_fn_trn(tfrec_dir, batchsize, input_height, input_width)
    test_dataset = input_fn_test(tfrec_dir, batchsize, input_height, input_width)

    '''
    Call backs
    '''
    tb_call = TensorBoard(log_dir=tboard)

    chkpt_call = ModelCheckpoint(filepath=os.path.join(chkpt_dir, 'f_model.h5'),
                                 monitor='val_accuracy',
                                 verbose=1,
                                 save_best_only=True)

    lr_scheduler_call = LearningRateScheduler(schedule=step_decay,
                                              verbose=1)

    callbacks_list = [tb_call, chkpt_call, lr_scheduler_call]

    '''
    Compile model
    Adam optimizer to change weights & biases
    Loss function is categorical crossentropy
    '''
    model.compile(optimizer=Adam(learning_rate=learnrate),
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    '''
    Training
    '''

    print('\n' + DIVIDER)
    print(' Training model with training set..')
    print(DIVIDER)

    # make folder for saving trained model checkpoint
    os.makedirs(chkpt_dir, exist_ok=True)

    # run training
    train_history = model.fit(train_dataset,
                              epochs=epochs,
                              steps_per_epoch=1800 // batchsize,
                              validation_data=test_dataset,
                              validation_steps=None,
                              callbacks=callbacks_list,
                              verbose=1)

    print(
        "\nTensorBoard can be opened with the command: tensorboard --logdir={dir} --host localhost --port 6006".format(
            dir=tboard))

    return


def run_main():
    # construct the argument parser and parse the arguments
    ap = argparse.ArgumentParser()
    ap.add_argument('-ih', '--input_height', type=int, default=512, help='Input image height in pixels.')
    ap.add_argument('-iw', '--input_width', type=int, default=512, help='Input image width in pixels.')
    ap.add_argument('-ic', '--input_chan', type=int, default=3, help='Number of input image channels.')
    ap.add_argument('-t', '--tfrec_dir', type=str, default='tfrecords',
                    help='Path of folder containing TFRecord files.')
    ap.add_argument('-b', '--batchsize', type=int, default=50,
                    help='Training batchsize. Must be an integer. Default is 50.')
    ap.add_argument('-e', '--epochs', type=int, default=10,
                    help='number of training epochs. Must be an integer. Default is 100.')
    ap.add_argument('-lr', '--learnrate', type=float, default=0.001,
                    help='optimizer learning rate. Must be floating-point value. Default is 0.0001')
    ap.add_argument('-cf', '--chkpt_dir', type=str, default='float_model',
                    help='Path and name of folder for storing Keras checkpoints. Default is float_model')
    ap.add_argument('-tb', '--tboard', type=str, default='tb_logs',
                    help='path to folder for saving TensorBoard data. Default is tb_logs.')
    args = ap.parse_args()

    print('\n' + DIVIDER)
    print('Keras version      : ', tf.keras.__version__)
    print('TensorFlow version : ', tf.__version__)
    print(sys.version)
    print(DIVIDER)
    print(' Command line options:')
    print('--input_height : ', args.input_height)
    print('--input_width  : ', args.input_width)
    print('--input_chan   : ', args.input_chan)
    print('--tfrec_dir    : ', args.tfrec_dir)
    print('--batchsize    : ', args.batchsize)
    print('--learnrate    : ', args.learnrate)
    print('--epochs       : ', args.epochs)
    print('--chkpt_dir    : ', args.chkpt_dir)
    print('--tboard       : ', args.tboard)
    print(DIVIDER)

    train(args.input_height, args.input_width, args.input_chan, args.tfrec_dir,
          args.batchsize, args.learnrate, args.epochs, args.chkpt_dir, args.tboard)


if __name__ == '__main__':
    run_main()
