import tensorflow as tf

class BlockingModel:
    def __init__(self):
        self.model = tf.keras.Sequential([
            # Embedding layer
            tf.keras.layers.Embedding(
                input_dim=90,    # like inputDim
                output_dim=32,   # like outputDim
                input_length=203,  # like inputLength
                mask_zero=True
            ),
            tf.keras.layers.Flatten(),  # flatten embedding output

            # Dense layers
            tf.keras.layers.Dense(512, activation='relu'),
            tf.keras.layers.Dropout(0.5),
            tf.keras.layers.Dense(256, activation='relu'),
            tf.keras.layers.Dropout(0.5),
            tf.keras.layers.Dense(128, activation='relu'),
            tf.keras.layers.Dropout(0.5),
            tf.keras.layers.Dense(1, activation='sigmoid')
        ])

        self.model.compile(
            optimizer=tf.keras.optimizers.RMSprop(0.01),
            loss='binary_crossentropy',
        )

    def predict(self, X: tf.Tensor) -> tf.Tensor:
        """Predict block probability"""
        return self.model(X, training=False)

    def train_batch(self, X: tf.Tensor, y: tf.Tensor):
        """Train on a single batch"""
        self.model.fit(X, y, batch_size=len(y), epochs=10, verbose=1)
