# First, make sure you have TensorFlow installed:
# pip install tensorflow

# Import the TensorFlow library, which includes the Keras API
import tensorflow as tf

# Define the model using the Sequential API, which allows you to create models layer-by-layer
model = tf.keras.Sequential()

# Add a convolutional layer with 32 filters, a kernel size of 3x3, 
# 'relu' activation, and specify the input shape that matches your preprocessed images
model.add(tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(64, 64, 3)))

# Add a max pooling layer to downsample the feature maps and reduce overfitting
model.add(tf.keras.layers.MaxPooling2D((2, 2)))

# Add another convolutional layer, this time with 64 filters
model.add(tf.keras.layers.Conv2D(64, (3, 3), activation='relu'))

# Add another max pooling layer
model.add(tf.keras.layers.MaxPooling2D((2, 2)))

# Flatten the 3D output to 1D for the dense layers
model.add(tf.keras.layers.Flatten())

# Add a dense layer with 64 units and 'relu' activation
model.add(tf.keras.layers.Dense(64, activation='relu'))

# Add the output layer with 1 unit and 'sigmoid' activation for binary classification
model.add(tf.keras.layers.Dense(1, activation='sigmoid'))

# Compile the model with an optimizer, a loss function for binary classification, 
# and specify that you want to track accuracy
model.compile(optimizer='adam',
              loss='binary_crossentropy',
              metrics=['accuracy'])

# Print out the model summary to see its structure
model.summary()

# Train the model using your preprocessed dataset
# X_train should be an array of your image data
# y_train should be an array of your labels (1 for red obstacles, 0 for no obstacles)
# validation_split indicates what fraction of the data to use for validation
# epochs indicates how many times to go through the entire training dataset
history = model.fit(X_train, y_train, epochs=10, validation_split=0.2)

# Evaluate the model's performance on the test dataset
# X_test should be your test image data
# y_test should be your test labels
test_loss, test_acc = model.evaluate(X_test, y_test)
print(f"Test accuracy: {test_acc}")

# Save the trained model for later use
model.save('red_obstacle_detection_model.h5')