"""
Resnet for speech commands classification.
"""

from typing import Optional

from art.estimators.classification import TensorFlowV2Classifier
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import losses
from tensorflow.keras.layers import Lambda


def get_spectrogram(audio):
    waveform = tf.convert_to_tensor(audio, dtype=tf.float32)
    spectrogram = tf.signal.stft(waveform, frame_length=255, frame_step=128)
    spectrogram = tf.abs(spectrogram)
    spectrogram = spectrogram[..., tf.newaxis]
    return spectrogram  # shape (124, 129, 1)


def make_audio_resnet(sequential=True, **kwargs) -> tf.keras.Model:

    inputs = keras.Input(shape=(16000,))
    spectrogram = Lambda(lambda audio: get_spectrogram(audio))(inputs)

    resnet = keras.applications.ResNet50(
        weights=None,
        input_tensor=spectrogram,
        classes=12,
        **kwargs,
    )

    model = keras.Model(resnet.inputs, resnet.outputs)
    if sequential:
        # ART's TensorFlowV2Classifier get_activations() requires a Sequential model
        model = keras.Sequential([model])

    model.compile(
        optimizer=tf.keras.optimizers.Adam(),
        loss=tf.keras.losses.SparseCategoricalCrossentropy(),
        metrics=["accuracy"],
    )

    return model


def get_art_model(
    model_kwargs: dict, wrapper_kwargs: dict, weights_path: Optional[str] = None
):

    model = make_audio_resnet(sequential=True, **model_kwargs)

    loss_object = losses.SparseCategoricalCrossentropy()

    @tf.function
    def train_step(model, samples, labels):
        with tf.GradientTape() as tape:
            predictions = model(samples, training=True)
            loss = loss_object(labels, predictions)
        gradients = tape.gradient(loss, model.trainable_variables)
        model.optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    art_classifier = TensorFlowV2Classifier(
        model=model,
        loss_object=loss_object,
        train_step=train_step,
        nb_classes=12,
        input_shape=(16000,),
        **wrapper_kwargs,
    )

    return art_classifier


def get_unwrapped_model(
    weights_path: str,
    **model_kwargs,
):
    # This is used for the explanatory model for the poisoning fairness metrics
    model = make_audio_resnet(sequential=False, **model_kwargs)
    model.load_weights(weights_path)

    return model
