"""Util for classification models."""

import os
import sys

import cv2
import numpy as np

from frigate.comms.embeddings_updater import EmbeddingsRequestEnum, EmbeddingsRequestor
from frigate.comms.inter_process import InterProcessRequestor
from frigate.const import CLIPS_DIR, MODEL_CACHE_DIR, UPDATE_MODEL_STATE
from frigate.types import ModelStatusTypesEnum
from frigate.util.process import FrigateProcess

BATCH_SIZE = 16
EPOCHS = 50
LEARNING_RATE = 0.001


def __generate_representative_dataset_factory(dataset_dir: str):
    def generate_representative_dataset():
        image_paths = []
        for root, dirs, files in os.walk(dataset_dir):
            for file in files:
                if file.lower().endswith((".jpg", ".jpeg", ".png")):
                    image_paths.append(os.path.join(root, file))

        for path in image_paths[:300]:
            img = cv2.imread(path)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = cv2.resize(img, (224, 224))
            img_array = np.array(img, dtype=np.float32) / 255.0
            img_array = img_array[None, ...]
            yield [img_array]

    return generate_representative_dataset


def __train_classification_model(model_name: str) -> bool:
    """Train a classification model."""

    # import in the function so that tensorflow is not initialized multiple times
    import tensorflow as tf
    from tensorflow.keras import layers, models, optimizers
    from tensorflow.keras.applications import MobileNetV2
    from tensorflow.keras.preprocessing.image import ImageDataGenerator

    dataset_dir = os.path.join(CLIPS_DIR, model_name, "dataset")
    model_dir = os.path.join(MODEL_CACHE_DIR, model_name)
    num_classes = len(
        [
            d
            for d in os.listdir(dataset_dir)
            if os.path.isdir(os.path.join(dataset_dir, d))
        ]
    )

    # TF and Keras are very loud with logging
    # we want to avoid these logs so we
    # temporarily redirect stdout / stderr
    original_stdout = sys.stdout
    original_stderr = sys.stderr
    sys.stdout = open(os.devnull, "w")
    sys.stderr = open(os.devnull, "w")

    # Start with imagenet base model with 35% of channels in each layer
    base_model = MobileNetV2(
        input_shape=(224, 224, 3),
        include_top=False,
        weights="imagenet",
        alpha=0.35,
    )
    base_model.trainable = False  # Freeze pre-trained layers

    model = models.Sequential(
        [
            base_model,
            layers.GlobalAveragePooling2D(),
            layers.Dense(128, activation="relu"),
            layers.Dropout(0.3),
            layers.Dense(num_classes, activation="softmax"),
        ]
    )

    model.compile(
        optimizer=optimizers.Adam(learning_rate=LEARNING_RATE),
        loss="categorical_crossentropy",
        metrics=["accuracy"],
    )

    # create training set
    datagen = ImageDataGenerator(rescale=1.0 / 255, validation_split=0.2)
    train_gen = datagen.flow_from_directory(
        dataset_dir,
        target_size=(224, 224),
        batch_size=BATCH_SIZE,
        class_mode="categorical",
        subset="training",
    )

    # write labelmap
    class_indices = train_gen.class_indices
    index_to_class = {v: k for k, v in class_indices.items()}
    sorted_classes = [index_to_class[i] for i in range(len(index_to_class))]
    with open(os.path.join(model_dir, "labelmap.txt"), "w") as f:
        for class_name in sorted_classes:
            f.write(f"{class_name}\n")

    # train the model
    model.fit(train_gen, epochs=EPOCHS, verbose=0)

    # convert model to tflite
    converter = tf.lite.TFLiteConverter.from_keras_model(model)
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    converter.representative_dataset = __generate_representative_dataset_factory(
        dataset_dir
    )
    converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
    converter.inference_input_type = tf.uint8
    converter.inference_output_type = tf.uint8
    tflite_model = converter.convert()

    # write model
    with open(os.path.join(model_dir, "model.tflite"), "wb") as f:
        f.write(tflite_model)

    # restore original stdout / stderr
    sys.stdout = original_stdout
    sys.stderr = original_stderr


@staticmethod
def kickoff_model_training(
    embeddingRequestor: EmbeddingsRequestor, model_name: str
) -> None:
    requestor = InterProcessRequestor()
    requestor.send_data(
        UPDATE_MODEL_STATE,
        {
            "model": model_name,
            "state": ModelStatusTypesEnum.training,
        },
    )

    # run training in sub process so that
    # tensorflow will free CPU / GPU memory
    # upon training completion
    training_process = FrigateProcess(
        target=__train_classification_model,
        name=f"model_training:{model_name}",
        args=(model_name,),
    )
    training_process.start()
    training_process.join()

    # reload model and mark training as complete
    embeddingRequestor.send_data(
        EmbeddingsRequestEnum.reload_classification_model.value,
        {"model_name": model_name},
    )
    requestor.send_data(
        UPDATE_MODEL_STATE,
        {
            "model": model_name,
            "state": ModelStatusTypesEnum.complete,
        },
    )
    requestor.stop()