import os
import tensorflow as tf
from huggingface_hub import hf_hub_download

@tf.keras.utils.register_keras_serializable()
class InceptionModel(tf.keras.Model):
    def __init__(
        self,
        dropout_rate: float,
        l2_reg: float,
        dense_units: int,
        *,
        name="InceptionV3",
        **kwargs,
    ):
        super().__init__(name=name, **kwargs)

        # Store for get_config
        self.dropout_rate = dropout_rate
        self.l2_reg = l2_reg
        self.dense_units = dense_units
        # Ensure L2 regularizer is correctly instantiated
        l2_regularizer = tf.keras.regularizers.L2(l2_reg)

        inception_base = tf.keras.applications.InceptionV3(
            include_top=False, # Keep False to add custom top layers
            input_shape=(256, 256, 3),
            pooling='max',
            weights=None
        )

        inputs = tf.keras.layers.Input(shape=(256, 256, 3))

        x = tf.keras.layers.Rescaling(1./255.)(inputs) # Scale pixel values to [0, 1]

        x = inception_base(x) # Pass through the InceptionV3 base
        x = tf.keras.layers.Dropout(dropout_rate)(x)
        x = tf.keras.layers.Dense(dense_units, activation="relu", kernel_regularizer=l2_regularizer)(x)
        outputs = tf.keras.layers.Dense(10, activation='softmax')(x) # Output probabilities for each class

        # Create the Keras Model
        self.net = tf.keras.Model(inputs=inputs, outputs=outputs)

    def call(self, inputs, training=False):
        return self.net(inputs, training=training)

    def get_config(self):
        config = super().get_config()
        config.update({
            "dropout_rate": self.dropout_rate,
            "l2_reg": self.l2_reg,
            "dense_units": self.dense_units,
        })
        return config

    @classmethod
    def from_config(cls, config):
        dropout_rate = config.pop("dropout_rate")
        l2_reg = config.pop("l2_reg")
        dense_units = config.pop("dense_units")
        return cls(dropout_rate, l2_reg, dense_units, **config)

def load_task_2_model(model_name='task_2_model_inception.keras'):
    model_path = hf_hub_download(
        repo_id="NickNam2710/predict_rice_diseases",
        filename=model_name,
        revision="main"
    )
    
    model = tf.keras.models.load_model(model_path)

    return model