import tensorflow as tf
from pathlib import Path
from kidney_classification.entity.config_entity import PrepareBaseModelConfig
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten, BatchNormalization, Dropout


class PrepareBaseModel:
    @staticmethod
    def prepare_full_model():
        VGG_model = Sequential()

        pretrained_model = tf.keras.applications.VGG16(
            include_top=False,
            input_shape=(150, 150, 3),
            pooling="max",
            classes=4,
            weights="imagenet",
        )

        VGG_model.add(pretrained_model)
        VGG_model.add(Flatten())
        VGG_model.add(Dense(512, activation="relu"))
        VGG_model.add(BatchNormalization())
        VGG_model.add(Dropout(0.5))

        VGG_model.add(Dense(4, activation="softmax"))
        pretrained_model.trainable = False

        VGG_model.compile(
            optimizer="adam",
            loss="sparse_categorical_crossentropy",
            metrics=["accuracy"],
        )

        return VGG_model

    def update_base_model(self, config: PrepareBaseModelConfig):
        full_model = self.prepare_full_model()

        full_model.summary()
        full_model.save(config.updated_base_model_path)

    @staticmethod
    def save_model(path: Path, model: tf.keras.Model):
        model.save(path)