# 1) Import libraries
# These are the tools I need for data handling, plotting graphs, building the CNN model, and checking how well it performs.
import os, random, itertools, numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras import datasets, layers, models, regularizers
from sklearn.metrics import classification_report, confusion_matrix

# 2) Set random seed
# This makes sure my results are consistent every time I run the code.
SEED = 42
tf.random.set_seed(SEED)
np.random.seed(SEED)
random.seed(SEED)

# 3) Load CIFAR-10 dataset
# using CIFAR-10, which has 60,000 small colour images in 10 categories like cat, truck, and airplane.
(x_train, y_train), (x_test, y_test) = datasets.cifar10.load_data()

# 4) Create validation set
# split the training data again to make a validation set. This helps me check how well the model is learning.
from sklearn.model_selection import train_test_split
x_train, x_val, y_train, y_val = train_test_split(
    x_train, y_train, test_size=0.1, random_state=SEED, stratify=y_train
)

# 5) Normalise the images
# scaled all pixel values to be between 0 and 1. This helps the model train faster and more smoothly.
x_train = x_train.astype("float32") / 255.0
x_val = x_val.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0

# 6) Define class names
# These are the labels for each category so I can display readable results later.
class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer',
               'dog', 'frog', 'horse', 'ship', 'truck']

# 7) Add data augmentation
# added random flips, rotations and zooms to the training images to help the model generalise better.
data_augmentation = tf.keras.Sequential([
    layers.RandomFlip("horizontal"),
    layers.RandomRotation(0.08),
    layers.RandomZoom(0.1)
], name="augment")

# 8) Build the CNN model
# This is the structure of my neural network. It has convolutional layers to detect patterns, pooling to reduce size, and dropout to prevent overfitting.
weight_decay = 1e-4
inputs = layers.Input(shape=(32, 32, 3))
x = data_augmentation(inputs)
x = layers.Conv2D(32, 3, padding='same', kernel_regularizer=regularizers.l2(weight_decay))(x)
x = layers.BatchNormalization()(x)
x = layers.ReLU()(x)
x = layers.MaxPooling2D()(x)

x = layers.Conv2D(64, 3, padding='same', kernel_regularizer=regularizers.l2(weight_decay))(x)
x = layers.BatchNormalization()(x)
x = layers.ReLU()(x)
x = layers.MaxPooling2D()(x)

x = layers.Conv2D(128, 3, padding='same', kernel_regularizer=regularizers.l2(weight_decay))(x)
x = layers.BatchNormalization()(x)
x = layers.ReLU()(x)
x = layers.GlobalAveragePooling2D()(x)
x = layers.Dropout(0.4)(x)

outputs = layers.Dense(10, activation='softmax')(x)
model = models.Model(inputs, outputs, name="cifar10_cnn")

# 9) Show model summary
# This prints out the structure of my model so I can see how many layers and parameters it has.
model.summary()

# 10) Compile the model
# chose the Adam optimiser and sparse categorical crossentropy loss. I’m measuring accuracy during training.
model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# 11) Add callbacks
# These help control training. EarlyStopping stops if the model stops improving. ReduceLROnPlateau lowers the learning rate if progress slows down.
early = tf.keras.callbacks.EarlyStopping(
    monitor='val_loss', patience=8, restore_best_weights=True
)
reduce = tf.keras.callbacks.ReduceLROnPlateau(
    monitor='val_loss', factor=0.5, patience=4, verbose=1, min_lr=1e-6
)

# 12) Train the model
# trained the model for up to 40 epochs using a batch size of 64. The validation set helps me monitor progress.
history = model.fit(
    x_train, y_train,
    validation_data=(x_val, y_val),
    epochs=40,
    batch_size=64,
    callbacks=[early, reduce],
    verbose=1
)

# 13) Evaluate the model
# tested the model on new images it hadn’t seen before. This gives me the final accuracy and loss.
test_loss, test_acc = model.evaluate(x_test, y_test, verbose=0)
print(f"Test accuracy: {test_acc:.4f} | Test loss: {test_loss:.4f}")

# 14) Plot accuracy and loss
# These graphs show how well the model learned over time. I can see if it overfitted or underfitted.
plt.figure()
plt.plot(history.history['accuracy'], label='train_acc')
plt.plot(history.history['val_accuracy'], label='val_acc')
plt.title('Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.show()

plt.figure()
plt.plot(history.history['loss'], label='train_loss')
plt.plot(history.history['val_loss'], label='val_loss')
plt.title('Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.show()

# 15) Generate predictions and confusion matrix
# checked how well the model did for each category. The confusion matrix shows where it got confused.
y_pred = model.predict(x_test, verbose=0)
y_pred_labels = np.argmax(y_pred, axis=1).reshape(-1)
y_true_labels = y_test.reshape(-1)

print("\nClassification Report:\n")
print(classification_report(y_true_labels, y_pred_labels, target_names=class_names, digits=4))

cm = confusion_matrix(y_true_labels, y_pred_labels)
plt.figure(figsize=(7, 6))
plt.imshow(cm, interpolation='nearest')
plt.title("Confusion Matrix")
plt.colorbar()
tick_marks = np.arange(len(class_names))
plt.xticks(tick_marks, class_names, rotation=45, ha='right')
plt.yticks(tick_marks, class_names)

thresh = cm.max() / 2.0
for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
    plt.text(j, i, format(cm[i, j], 'd'),
             horizontalalignment="center",
             color="white" if cm[i, j] > thresh else "black", fontsize=8)

plt.tight_layout()
plt.ylabel('True label')
plt.xlabel('Predicted label')
plt.show()

# 16) Show sample predictions
# displayed 16 test images with their true and predicted labels. Green means correct, red means incorrect.
def plot_samples(images, true_labs, pred_labs, n=16):
    idx = np.random.choice(len(images), n, replace=False)
    images = images[idx]
    true_labs = true_labs[idx]
    pred_labs = pred_labs[idx]
    cols = 8
    rows = n // cols
    plt.figure(figsize=(16, 4 * rows))
    for i in range(n):
        plt.subplot(rows, cols, i + 1)
        plt.imshow(images[i])
        t = class_names[int(true_labs[i])]
        p = class_names[int(pred_labs[i])]
        title = f"T: {t} | P: {p}"
        colour = 'green' if t == p else 'red'
        plt.title(title, color=colour, fontsize=9)
        plt.axis('off')
    plt.tight_layout()
    plt.show()

plot_samples(x_test, y_true_labels, y_pred_labels, n=16)

# 17) Save the model
# saved the trained model so I can reuse it later without retraining.
model.save("cifar10_cnn.h5")
print("Saved model to cifar10_cnn.h5")