import numpy as np
import tensorflow as tf
from tensorflow.keras import datasets, layers, models

# Function to convert images to grayscale
def rgb_to_grayscale(images):
    # Convert images to grayscale using the NTSC formula
    return np.dot(images[...,:3], [0.2989, 0.5870, 0.1140])

# Load CIFAR-10 dataset
(train_images, train_labels), (test_images, test_labels) = datasets.cifar10.load_data()

# Convert the images to grayscale
train_images_grayscale = rgb_to_grayscale(train_images).reshape(-1, 32, 32, 1)
test_images_grayscale = rgb_to_grayscale(test_images).reshape(-1, 32, 32, 1)

# Normalize pixel values to be between 0 and 1
train_images_grayscale, test_images_grayscale = train_images_grayscale / 255.0, test_images_grayscale / 255.0

# Define the CNN model
model = models.Sequential()
# Sliding Window Convolution Layer
model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 1)))
# Max Pooling Layer
model.add(layers.MaxPooling2D((2, 2)))
# Second Convolution Layer
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
# Second Max Pooling Layer
model.add(layers.MaxPooling2D((2, 2)))
# Third Convolution Layer
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
# Flattening Layer
model.add(layers.Flatten())
# Fully Connected Layer
model.add(layers.Dense(64, activation='relu'))
# Softmax Activation Layer
model.add(layers.Dense(10, activation='softmax'))

# Compile the model
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(filepath='weights_epoch_{epoch:02d}.h5', save_weights_only=False, verbose=1)


# Train the model
history = model.fit(train_images_grayscale, train_labels, epochs=50,
                    validation_data=(test_images_grayscale, test_labels),callbacks=[checkpoint_callback])


# Print a success message
print("CNN model for CIFAR-10 dataset with grayscale images has been successfully created and trained.")