import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.datasets import cifar10
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
from tensorflow.keras.utils import to_categorical

# -----------------------------------
# 1. Load and Preprocess CIFAR-10 Data
# -----------------------------------
(train_images, train_labels), (test_images, test_labels) = cifar10.load_data()

# Normalize pixel values to range [0, 1]
train_images = train_images.astype("float32") / 255.0
test_images = test_images.astype("float32") / 255.0

# One-hot encode class labels (10 classes)
train_labels_cat = to_categorical(train_labels, 10)
test_labels_cat = to_categorical(test_labels, 10)

# -----------------------------------
# 2. Build CNN Model
# -----------------------------------
cnn_model = Sequential([
    Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3)),
    MaxPooling2D(pool_size=(2, 2)),

    Conv2D(64, (3, 3), activation='relu'),
    MaxPooling2D(pool_size=(2, 2)),

    Conv2D(128, (3, 3), activation='relu'),
    Flatten(),

    Dense(128, activation='relu'),
    Dropout(0.5),
    Dense(10, activation='softmax')  # 10 output classes
])

# -----------------------------------
# 3. Compile the Model
# -----------------------------------
cnn_model.compile(
    optimizer='adam',
    loss='categorical_crossentropy',
    metrics=['accuracy']
)

# -----------------------------------
# 4. Train the Model
# -----------------------------------
history = cnn_model.fit(
    train_images, train_labels_cat,
    epochs=10,
    batch_size=64,
    validation_split=0.1,
    verbose=2
)

# -----------------------------------
# 5. Evaluate on Test Set
# -----------------------------------
test_loss, test_accuracy = cnn_model.evaluate(test_images, test_labels_cat, verbose=0)
print(f"✅ Test Accuracy: {test_accuracy * 100:.2f}%")

# -----------------------------------
# 6. Plot Accuracy Trends
# -----------------------------------
plt.plot(history.history['accuracy'], label='Train Accuracy')
plt.plot(history.history['val_accuracy'], label='Validation Accuracy')
plt.title('CNN Accuracy on CIFAR-10')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()