# ============================================================
# Machine Learning Task Two
# CIFAR-10 Image Classification using a CNN
# ============================================================

# ------------------------------------------------------------
# STEP 0: Install TensorFlow if it is missing
# ------------------------------------------------------------
# This section checks whether TensorFlow is installed.
# If TensorFlow is missing, the code installs the CPU version.
# The CPU version is fine for this coursework and avoids GPU issues.

try:
    import tensorflow as tf
    print("TensorFlow is already installed.")
    print("TensorFlow version:", tf.__version__)

except ModuleNotFoundError:
    import sys
    import subprocess

    print("TensorFlow not found. Installing tensorflow-cpu now...")

    subprocess.check_call([
        sys.executable,
        "-m",
        "pip",
        "install",
        "tensorflow-cpu"
    ])

    import tensorflow as tf
    print("TensorFlow installed successfully.")
    print("TensorFlow version:", tf.__version__)


# ------------------------------------------------------------
# STEP 1: Import required libraries
# ------------------------------------------------------------

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
from tensorflow.keras.optimizers import Adam

from sklearn.metrics import classification_report, confusion_matrix, ConfusionMatrixDisplay


# ------------------------------------------------------------
# STEP 2: Load the CIFAR-10 dataset
# ------------------------------------------------------------
# CIFAR-10 contains 60,000 colour images.
# 50,000 images are used for training.
# 10,000 images are used for testing.
# Each image is 32 x 32 pixels with 3 colour channels.

(x_train, y_train), (x_test, y_test) = cifar10.load_data()

print("Training image shape:", x_train.shape)
print("Training label shape:", y_train.shape)
print("Testing image shape:", x_test.shape)
print("Testing label shape:", y_test.shape)


# ------------------------------------------------------------
# STEP 3: Define the CIFAR-10 class names
# ------------------------------------------------------------
# The dataset labels are numbers from 0 to 9.
# This list converts those numbers into readable class names.

class_names = [
    "airplane",
    "automobile",
    "bird",
    "cat",
    "deer",
    "dog",
    "frog",
    "horse",
    "ship",
    "truck"
]


# ------------------------------------------------------------
# STEP 4: Display sample images from the dataset
# ------------------------------------------------------------
# This helps prove that the correct CIFAR-10 dataset has loaded.

plt.figure(figsize=(10, 5))

for i in range(10):
    plt.subplot(2, 5, i + 1)
    plt.imshow(x_train[i])
    plt.title(class_names[int(y_train[i])])
    plt.axis("off")

plt.suptitle("Sample Images from the CIFAR-10 Dataset")
plt.show()


# ------------------------------------------------------------
# STEP 5: Data preprocessing
# ------------------------------------------------------------
# Neural networks train better when input values are small.
# Original pixel values range from 0 to 255.
# Dividing by 255 changes the range to 0 to 1.
# This is called normalisation.

x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0

print("Minimum pixel value after normalisation:", x_train.min())
print("Maximum pixel value after normalisation:", x_train.max())


# ------------------------------------------------------------
# STEP 6: One-hot encode the labels
# ------------------------------------------------------------
# CIFAR-10 labels are stored as numbers.
# Example: cat may be label 3.
# For categorical cross-entropy, labels are converted into
# one-hot encoded vectors.
#
# Example:
# class 3 becomes [0, 0, 0, 1, 0, 0, 0, 0, 0, 0]

y_train_cat = to_categorical(y_train, 10)
y_test_cat = to_categorical(y_test, 10)

print("Example original label:", y_train[0])
print("Example one-hot encoded label:", y_train_cat[0])


# ------------------------------------------------------------
# STEP 7: Build the CNN model
# ------------------------------------------------------------
# The model is kept simple to match the assignment requirement.
#
# Conv2D:
# Extracts image features such as edges, shapes, and textures.
#
# MaxPooling2D:
# Reduces image feature size and lowers computational cost.
#
# Flatten:
# Converts the 2D feature maps into a 1D vector.
#
# Dense:
# Performs classification based on extracted features.
#
# Dropout:
# Reduces overfitting by randomly turning off neurons during training.
#
# Softmax:
# Produces probabilities for the 10 CIFAR-10 classes.

model = Sequential()

# First convolution and pooling block
model.add(Conv2D(
    filters=32,
    kernel_size=(3, 3),
    activation="relu",
    input_shape=(32, 32, 3)
))
model.add(MaxPooling2D(pool_size=(2, 2)))

# Second convolution and pooling block
model.add(Conv2D(
    filters=64,
    kernel_size=(3, 3),
    activation="relu"
))
model.add(MaxPooling2D(pool_size=(2, 2)))

# Convert feature maps into one long vector
model.add(Flatten())

# Fully connected hidden layer
model.add(Dense(128, activation="relu"))

# Dropout layer to reduce overfitting
model.add(Dropout(0.5))

# Output layer for 10 classes
model.add(Dense(10, activation="softmax"))


# ------------------------------------------------------------
# STEP 8: Compile the model
# ------------------------------------------------------------
# Adam optimiser:
# Used because it is efficient and adapts the learning rate.
#
# Categorical cross-entropy:
# Used because this is a multi-class classification problem.
#
# Accuracy:
# Used to measure the proportion of correct predictions.

model.compile(
    optimizer=Adam(),
    loss="categorical_crossentropy",
    metrics=["accuracy"]
)

print("\nCNN Model Summary:")
model.summary()


# ------------------------------------------------------------
# STEP 9: Train the model
# ------------------------------------------------------------
# Epochs:
# Number of times the model sees the full training dataset.
#
# Batch size:
# Number of images processed before the model updates its weights.
#
# Validation split:
# 20% of the training data is used to monitor validation performance.

history = model.fit(
    x_train,
    y_train_cat,
    epochs=10,
    batch_size=32,
    validation_split=0.2,
    verbose=1
)


# ------------------------------------------------------------
# STEP 10: Evaluate the model on the test set
# ------------------------------------------------------------
# The test set is unseen data.
# This gives a better idea of real model performance.

test_loss, test_accuracy = model.evaluate(
    x_test,
    y_test_cat,
    verbose=0
)

print("\nFinal Test Results")
print("------------------")
print(f"Test Accuracy: {test_accuracy:.4f}")
print(f"Test Loss: {test_loss:.4f}")


# ------------------------------------------------------------
# STEP 11: Plot training and validation accuracy
# ------------------------------------------------------------
# This graph shows whether the model improved during training.
# It also helps identify overfitting.

plt.figure(figsize=(8, 5))
plt.plot(history.history["accuracy"], label="Training Accuracy")
plt.plot(history.history["val_accuracy"], label="Validation Accuracy")
plt.title("CNN Training and Validation Accuracy")
plt.xlabel("Epoch")
plt.ylabel("Accuracy")
plt.legend()
plt.grid(True)
plt.show()


# ------------------------------------------------------------
# STEP 12: Plot training and validation loss
# ------------------------------------------------------------
# Loss measures prediction error.
# A decreasing loss usually means the model is learning.

plt.figure(figsize=(8, 5))
plt.plot(history.history["loss"], label="Training Loss")
plt.plot(history.history["val_loss"], label="Validation Loss")
plt.title("CNN Training and Validation Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.grid(True)
plt.show()


# ------------------------------------------------------------
# STEP 13: Make predictions
# ------------------------------------------------------------
# The model outputs probabilities for each class.
# np.argmax selects the class with the highest probability.

y_pred_probs = model.predict(x_test)
y_pred = np.argmax(y_pred_probs, axis=1)
y_true = y_test.flatten()


# ------------------------------------------------------------
# STEP 14: Display classification report
# ------------------------------------------------------------
# Precision:
# How many predicted positives were correct.
#
# Recall:
# How many actual positives were found.
#
# F1-score:
# Balance between precision and recall.

print("\nClassification Report")
print("---------------------")
print(classification_report(
    y_true,
    y_pred,
    target_names=class_names
))


# ------------------------------------------------------------
# STEP 15: Create confusion matrix
# ------------------------------------------------------------
# The confusion matrix shows which classes were predicted correctly
# and which classes were confused with others.

cm = confusion_matrix(y_true, y_pred)

plt.figure(figsize=(10, 8))
disp = ConfusionMatrixDisplay(
    confusion_matrix=cm,
    display_labels=class_names
)

disp.plot(cmap="Blues", xticks_rotation=45)
plt.title("CNN Confusion Matrix for CIFAR-10")
plt.show()


# ------------------------------------------------------------
# STEP 16: Display some example predictions
# ------------------------------------------------------------
# This gives a visual check of how the model performs.

plt.figure(figsize=(12, 8))

for i in range(15):
    plt.subplot(3, 5, i + 1)
    plt.imshow(x_test[i])
    
    true_label = class_names[y_true[i]]
    predicted_label = class_names[y_pred[i]]
    
    plt.title(f"True: {true_label}\nPred: {predicted_label}")
    plt.axis("off")

plt.suptitle("Example CIFAR-10 Predictions")
plt.tight_layout()
plt.show()