"""
图像分割模型构建组件。

这个文件承载 notebook 中的编码器/解码器分割网络，以及监督学习 Pipeline 使用的
模型构建器。
"""

from dataclasses import dataclass

import keras
from keras.layers import Conv2D, Conv2DTranspose, Rescaling

from deep_learning.models.spec import ModelArtifact, SupervisedModelBuilder


@dataclass
class SegmentationModelBuilder(SupervisedModelBuilder):
    image_size: tuple[int, int]
    num_classes: int
    model_filters: tuple[int, ...]

    def build_training_artifact(self) -> ModelArtifact:
        inputs = keras.Input(shape=self.image_size + (3,))
        x = Rescaling(1.0 / 255)(inputs)

        for filter_count in self.model_filters:
            x = Conv2D(filter_count, 3, strides=2, activation="relu", padding="same")(x)
            x = Conv2D(filter_count, 3, activation="relu", padding="same")(x)

        for filter_count in reversed(self.model_filters):
            x = Conv2DTranspose(filter_count, 3, activation="relu", padding="same")(x)
            x = Conv2DTranspose(filter_count, 3, strides=2, activation="relu", padding="same")(x)

        outputs = Conv2D(self.num_classes, 3, activation="softmax", padding="same")(x)
        model = keras.Model(inputs, outputs, name="segmentation")
        return ModelArtifact(model=model)

    def compile_training_model(self, model: keras.Model) -> None:
        foreground_iou = keras.metrics.IoU(
            num_classes=self.num_classes,
            target_class_ids=(0,),
            name="foreground_iou",
            sparse_y_true=True,
            sparse_y_pred=False
        )
        model.compile(
            optimizer="adam",
            loss="sparse_categorical_crossentropy",
            metrics=[foreground_iou]
        )