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


def mish(inputs):
    """Mish activation function

    Args:
        inputs (tf.Tensor): input tensor

    Returns:
        tf.Tensor: output tensor
    """
    x = tf.nn.softplus(inputs)
    x = tf.nn.tanh(x)
    x = tf.multiply(x, inputs)
    return x


def srm_24_16(input_shape, classes):
    x = inputs = tf.keras.layers.Input(input_shape)
    x = tf.keras.layers.Dense(128, activation=mish)(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    x = tf.keras.layers.Masking(mask_value=0.0)(x)
    x = tf.keras.layers.GRU(24, return_sequences=True, activation='tanh')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    x = tf.keras.layers.GRU(16, return_sequences=False, activation='tanh')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    x = tf.keras.layers.Dense(classes, activation='softmax')(x)
    return tf.keras.Model(inputs, x)


def srm_32_24(input_shape, classes):
    x = inputs = tf.keras.layers.Input(input_shape)
    x = tf.keras.layers.Dense(128, activation=mish)(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    x = tf.keras.layers.Masking(mask_value=0.0)(x)
    x = tf.keras.layers.GRU(32, return_sequences=True, activation='tanh')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    x = tf.keras.layers.GRU(24, return_sequences=False, activation='tanh')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    x = tf.keras.layers.Dense(classes, activation='softmax')(x)
    return tf.keras.Model(inputs, x)


def srm_40_40(input_shape, classes):
    x = inputs = tf.keras.layers.Input(input_shape)
    x = tf.keras.layers.Dense(128, activation=mish)(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    x = tf.keras.layers.Masking(mask_value=0.0)(x)
    x = tf.keras.layers.GRU(40, return_sequences=True, activation='tanh')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    x = tf.keras.layers.GRU(40, return_sequences=False, activation='tanh')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    x = tf.keras.layers.Dense(classes, activation='softmax')(x)
    return tf.keras.Model(inputs, x)


def srm_48_32_32(input_shape, classes):
    x = inputs = tf.keras.layers.Input(input_shape)
    x = tf.keras.layers.Dense(128, activation=mish)(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    x = tf.keras.layers.Masking(mask_value=0.0)(x)
    x = tf.keras.layers.GRU(48, return_sequences=True, activation='tanh')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    x = tf.keras.layers.GRU(32, return_sequences=True, activation='tanh')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    x = tf.keras.layers.GRU(32, return_sequences=False, activation='tanh')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    x = tf.keras.layers.Dense(classes, activation='softmax')(x)
    return tf.keras.Model(inputs, x)


def SignRecognitionModel(input_shape, classes):
    """Sign langauge classification model

    Args:
        input_shape (tuple): input feature shape
        classes (int): number of classes, one of (100, 300, 1000, 2000)

    Returns:
        tf.keras.Model: Classification model
    """
    model_dict = {
        100: srm_24_16,
        300: srm_32_24,
        1000: srm_40_40,
        2000: srm_48_32_32,
    }
    return model_dict[classes](input_shape, classes)
