"""
Train WLASL classification models
"""
import tensorflow as tf

from modules.dataset_utils import init_dataset
from modules.model import SignRecognitionModel
from modules.config import CFG


def train_classifier(n_classes):
    """Train classification model

    Args:
        n_classes (int): number of classes - one of (100, 300, 1000, 2000)
    """
    # get training and validation datasets
    train_data = init_dataset(
        CFG.WLASL_TFRECORDS_PATH[n_classes]['train'],
        CFG.WLASL_PARAMS.BATCH_SIZE,
        n_classes,
        CFG.WLASL_PARAMS.SEQUENCE_LEN,
        augment=True,
        shuffle=True
    )
    val_data = init_dataset(
        CFG.WLASL_TFRECORDS_PATH[n_classes]['val'],
        CFG.WLASL_PARAMS.BATCH_SIZE,
        n_classes,
        CFG.WLASL_PARAMS.SEQUENCE_LEN,
    )
    # get model
    model = SignRecognitionModel((CFG.WLASL_PARAMS.SEQUENCE_LEN, 2*21*2+1*98*2), n_classes)
    model.summary()
    # cross-entropy loss function
    loss = tf.keras.losses.CategoricalCrossentropy()
    # callbacks
    callbacks = [
            tf.keras.callbacks.TensorBoard(log_dir=f'logs/sign_recognition_{n_classes}'),
            tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=200, verbose=True, restore_best_weights=True, min_delta=0.01),
            tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', patience=100)
        ]
    # optimizer
    optimizer = tf.keras.optimizers.Nadam(learning_rate=CFG.WLASL_PARAMS.LEARNING_RATE)
    # compile, train and save weights
    model.compile(loss=loss, optimizer=optimizer, metrics=['accuracy'])
    model.fit(train_data, epochs=CFG.WLASL_PARAMS.EPOCHS, callbacks=callbacks, validation_data=val_data)
    model.save_weights(CFG.WLASL_WEIGHTS[n_classes])
    # get testing dataset and evaluate on top-k accuracy
    test_data = init_dataset(
        CFG.WLASL_TFRECORDS_PATH[n_classes]['test'],
        1,
        n_classes
    )
    model.compile(metrics=[
        tf.keras.metrics.TopKCategoricalAccuracy(k=1, name='top_1_accuracy'),
        tf.keras.metrics.TopKCategoricalAccuracy(k=5, name='top_5_accuracy'),
        tf.keras.metrics.TopKCategoricalAccuracy(k=10, name='top_10_accuracy'),
    ])
    model.evaluate(test_data)


if __name__ == '__main__':
    train_classifier(100)
    train_classifier(300)
    train_classifier(1000)
    train_classifier(2000)
