import tensorflow as tf

# Custom conv2d Layer
class Conv2DBatchNMaxP(tf.keras.layers.Layer):
    # TAMBAHKAN **kwargs di sini
    def __init__(self, inp1, inp2, **kwargs):
        super(Conv2DBatchNMaxP, self).__init__(**kwargs) # Teruskan kwargs
        self.inp1 = inp1
        self.inp2 = inp2
        self.conv2d_main = tf.keras.layers.Conv2D(32, (inp1, inp2), padding='same', activation='relu')
        self.bn_main = tf.keras.layers.BatchNormalization()
        self.maxpool2d_main = tf.keras.layers.MaxPool2D((2,2))
        self.maxpool2d_shortcut = tf.keras.layers.MaxPool2D((2,2))

    def call(self, inputs):
        x = self.conv2d_main(inputs)
        x = self.bn_main(x)
        x = self.maxpool2d_main(x)
        shortcut = self.maxpool2d_shortcut(inputs)
        return x + shortcut

    # Tambahkan get_config agar layer bisa disimpan/dimuat dengan benar
    def get_config(self):
        config = super(Conv2DBatchNMaxP, self).get_config()
        config.update({
            "inp1": self.inp1,
            "inp2": self.inp2,
        })
        return config


# Model Subclassing
class Conv2DModel(tf.keras.Model):
    # TAMBAHKAN **kwargs di sini untuk menangani parameter 'trainable', 'dtype', dll.
    def __init__(self, **kwargs):
        super(Conv2DModel, self).__init__(**kwargs) # Teruskan kwargs
        self.ConvLayer1a = tf.keras.layers.Conv2D(32, (3,3), padding='same', activation='relu', input_shape=(150,150,1))
        self.ConvLayer1b = tf.keras.layers.BatchNormalization()
        self.ConvLayer1c = tf.keras.layers.MaxPool2D((2,2))
        self.ConvLayer2 = Conv2DBatchNMaxP(3, 3)
        self.ConvLayer3 = Conv2DBatchNMaxP(3, 3)
        self.ConvLayer4 = Conv2DBatchNMaxP(3, 3)
        self.flatten = tf.keras.layers.Flatten()
        
        # Dense layers
        self.dense1 = tf.keras.layers.Dense(256, activation='relu')
        self.dropout1 = tf.keras.layers.Dropout(0.5)
        self.dense2 = tf.keras.layers.Dense(128, activation='relu')
        self.dropout2 = tf.keras.layers.Dropout(0.5)
        self.dense3 = tf.keras.layers.Dense(128, activation='relu')
        self.dropout3 = tf.keras.layers.Dropout(0.5)
        self.dense4 = tf.keras.layers.Dense(128, activation='relu')
        self.dropout4 = tf.keras.layers.Dropout(0.5)
        self.dense5 = tf.keras.layers.Dense(128, activation='relu')
        self.dropout5 = tf.keras.layers.Dropout(0.5)
        self.dense6 = tf.keras.layers.Dense(128, activation='relu')
        self.dropout6 = tf.keras.layers.Dropout(0.5)
        self.dense7 = tf.keras.layers.Dense(128, activation='relu')
        self.dropout7 = tf.keras.layers.Dropout(0.5)
        self.dense8 = tf.keras.layers.Dense(128, activation='relu')
        self.dropout8 = tf.keras.layers.Dropout(0.5)
        
        # Output layer (Sesuai jumlah kelas dataset saat training!)
        self.dense9 = tf.keras.layers.Dense(31, activation='softmax') 

    def call(self, inputs):
        x = self.ConvLayer1a(inputs)
        x = self.ConvLayer1b(x)
        x = self.ConvLayer1c(x)
        x = self.ConvLayer2(x)
        x = self.ConvLayer3(x)
        x = self.ConvLayer4(x)
        x = self.flatten(x)
        x = self.dense1(x)
        x = self.dropout1(x)
        x = self.dense2(x)
        x = self.dropout2(x)
        x = self.dense3(x)
        x = self.dropout3(x)
        x = self.dense4(x)
        x = self.dropout4(x)
        x = self.dense5(x)
        x = self.dropout5(x)
        x = self.dense6(x)
        x = self.dropout6(x)
        x = self.dense7(x)
        x = self.dropout7(x)
        x = self.dense8(x)
        x = self.dropout8(x)
        return self.dense9(x)