from dataclasses import dataclass
from pathlib import Path

import keras
import keras_hub
from keras import layers

from .losses import box_loss
from deep_learning.models.spec import ModelArtifact, SupervisedModelBuilder


def build_yolo_preprocessor(image_size: int = 448):
    inputs = keras.Input(shape=(None, None, 3), dtype="uint8")
    x = layers.Resizing(
        image_size,
        image_size,
        interpolation="bicubic",
        crop_to_aspect_ratio=True
    )(inputs)
    x = layers.Rescaling(
        scale=[0.017124753831663668, 0.01750700280112045, 0.017429193899782133],
        offset=[-2.1179039301310043, -2.0357142857142856, -1.8044444444444445]
    )(x)
    return keras.Model(inputs, x, name="yolo_preprocessor")


@dataclass
class YoloModelBuilder(SupervisedModelBuilder):
    image_size: int
    grid_size: int
    num_labels: int
    backbone_preset: str = "resnet_50_imagenet"

    def build_training_artifact(self) -> ModelArtifact:
        backbone = keras_hub.models.Backbone.from_preset(self.backbone_preset)
        inputs = keras.Input(shape=(self.image_size, self.image_size, 3))
        x = backbone(inputs)
        x = layers.Conv2D(512, (3, 3), strides=(2, 2))(x)
        x = layers.Flatten()(x)
        x = layers.Dense(2048, activation="relu", kernel_initializer="glorot_normal")(x)
        x = layers.Dropout(0.5)(x)
        x = layers.Dense(self.grid_size * self.grid_size * (self.num_labels + 5))(x)
        x = layers.Reshape((self.grid_size, self.grid_size, self.num_labels + 5))(x)
        box_predictions = x[..., :5]
        class_predictions = layers.Activation("softmax")(x[..., 5:])
        outputs = {"box": box_predictions, "class": class_predictions}
        model = keras.Model(inputs, outputs, name="yolo")
        return ModelArtifact(model=model)

    def convert_to_inference_artifact(
        self,
        training_artifact: ModelArtifact
    ) -> ModelArtifact:
        inputs = keras.Input(shape=(None, None, 3), dtype="uint8")
        preprocessor = build_yolo_preprocessor(image_size=self.image_size)
        x = preprocessor(inputs)
        outputs = training_artifact.model(x)
        inference_model = keras.Model(inputs, outputs, name="yolo_inference")
        return ModelArtifact(model=inference_model)

    def load_inference_artifact(self, model_path: Path) -> ModelArtifact:
        model = keras.models.load_model(
            str(model_path),
            custom_objects=self._custom_objects()
        )
        return ModelArtifact(model=model)

    def _custom_objects(self) -> dict:
        return {
            "box_loss": box_loss
        }

    def compile_training_model(self, model: keras.Model) -> None:
        model.compile(
            optimizer=keras.optimizers.Adam(2e-4),
            loss={
                "box": box_loss,
                "class": "sparse_categorical_crossentropy"
            }
        )