import tensorflow as tf
from tensorflow import keras
from ultralytics import YOLO
from official.projects.movinet.modeling.movinet import Movinet
from official.projects.movinet.modeling.movinet_model import MovinetClassifier

from configuration import Config

class AttentionDenseClassifierHead(keras.layers.Layer):
  def __init__(self, attention_heads, dense_units, dropout_rate=0.2, **kwargs):
    super().__init__(**kwargs)
    self.attention = keras.layers.MultiHeadAttention(num_heads=attention_heads, key_dim=1)
    self.normalization = keras.layers.LayerNormalization(epsilon=1e-6)
    self.dropout = keras.layers.Dropout(dropout_rate)
    self.dense = keras.layers.Dense(dense_units, activation='softmax')

  def call(self, x, training):
    y = tf.expand_dims(x, -1)
    y = self.attention(query=y, key=y, value=y)
    y = tf.squeeze(y, axis=-1)
    y = self.dropout(y, training=training)
    y = self.normalization(x + y*0.01)
    y = self.dense(y)
    return y

def build_movinet(output_size, config: Config):
  model = MovinetClassifier(
    backbone=Movinet(model_id=config.model_id),
    num_classes=output_size)
  model.build(config.input_shape)
  return model

def build_classifier_head(input_size, config: Config):
  inputs = keras.Input(shape=(input_size,))
  classifier = AttentionDenseClassifierHead(2, config.num_classes)(inputs)
  model = keras.Model(inputs=inputs, outputs=classifier)
  return model

def build_model(movinet, classifier_head):
  return keras.models.Sequential([movinet, classifier_head])

def load_classifier(config: Config):
  movinet = build_movinet(600, config)
  classifier_head = build_classifier_head(600, config)
  model = build_model(movinet, classifier_head)
  model.load_weights(config.classifier_path)
  return model

def load_detector(config: Config):
  return YOLO(config.detector_path)

def compile_classifier(model, config: Config):
  optimizer = keras.optimizers.Adam(learning_rate=config.learning_rate)
  model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])