import tensorflow as tf
from tensorflow.keras import layers, Model


def separable_block(x, filters, dropout_rate=0.25):
    x = layers.SeparableConv2D(filters, 3, padding="same", use_bias=False)(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)

    x = layers.SeparableConv2D(filters, 3, padding="same", use_bias=False)(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)

    x = layers.MaxPooling2D()(x)
    x = layers.SpatialDropout2D(dropout_rate)(x)

    return x



# Model Definition
def build_sara_tf_model(input_shape=(150, 150, 3), num_classes=6):

    inputs = tf.keras.Input(shape=input_shape)

   
    x = layers.Conv2D(32, 3, padding="same", use_bias=False)(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)

    
    x = separable_block(x, 64, 0.25)
    x = separable_block(x, 128, 0.25)
    x = separable_block(x, 256, 0.30)

    
    x = layers.GlobalAveragePooling2D()(x)

    x = layers.Dense(128, use_bias=False)(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)
    x = layers.Dropout(0.5)(x)

    outputs = layers.Dense(num_classes, activation="softmax")(x)

    model = Model(inputs, outputs, name="SaraCNN_TF")

    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
        loss="categorical_crossentropy",
        metrics=["accuracy"],
    )

    return model