# # CIFAR-10

# ## Ładowanie zbioru danych

# In[ ]:
import sys

import tensorflow as tf

# Check if GPU is available
print(tf.config.list_physical_devices("GPU"))
if tf.config.list_physical_devices("GPU"):
    print("GPU is available")
else:
    print("GPU is not available")
    sys.exit()


import itertools

import matplotlib.pyplot as plt
import numpy as np
from keras.datasets import cifar10
from keras.layers import (
    Activation,
    BatchNormalization,
    Conv2D,
    Dense,
    Dropout,
    Flatten,
    MaxPooling2D,
)
from keras.models import Sequential
from keras.utils import to_categorical
from sklearn.metrics import confusion_matrix
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.preprocessing.image import ImageDataGenerator


def plot_confusion_matrix(
    cm, classes, normalize=False, title="Confusion matrix", cmap=plt.cm.Blues
):
    """This function prints and plots the confusion matrix.
    Normalization can be applied by setting `normalize=True`.
    """
    if normalize:
        cm = cm.astype("float") / cm.sum(axis=1)[:, np.newaxis]
        print("Normalized confusion matrix")
    else:
        print("Confusion matrix, without normalization")

    print(cm)

    plt.imshow(cm, interpolation="nearest", cmap=cmap)
    plt.title(title)
    plt.colorbar()
    tick_marks = np.arange(len(classes))
    plt.xticks(tick_marks, classes, rotation=45)
    plt.yticks(tick_marks, classes)

    fmt = ".2f" if normalize else "d"
    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], fmt),
            horizontalalignment="center",
            color="white" if cm[i, j] > thresh else "black",
        )

    plt.ylabel("True label")
    plt.xlabel("Predicted label")
    plt.tight_layout()


(X_train, y_train), (X_test, y_test) = cifar10.load_data()

X_train = X_train.astype("float32")  # change integers to 32-bit floating point numbers
X_test = X_test.astype("float32")

X_train /= (
    255  # normalize each value for each pixel for the entire vector for each input
)
X_test /= 255

y_train = y_train.reshape((1, -1))[0]
y_test = y_test.reshape((1, -1))[0]

print("Training matrix shape", X_train.shape, y_train.shape)
print("Testing matrix shape", X_test.shape, y_test.shape)

# one-hot format classes

nb_classes = 10

Y_train = to_categorical(y_train, nb_classes)
Y_test = to_categorical(y_test, nb_classes)

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


# ## Podgląd zbioru treningowego

# In[ ]:


for i in range(10):
    img_batch = X_train[y_train == i][0:10]
    img_batch = np.reshape(
        img_batch,
        (
            img_batch.shape[0] * img_batch.shape[1],
            img_batch.shape[2],
            img_batch.shape[3],
        ),
    )
    if i > 0:
        img = np.concatenate([img, img_batch], axis=1)
    else:
        img = img_batch
plt.figure(figsize=(10, 20))
plt.axis("off")
plt.imshow(img, cmap="gray")


# ## Przygotowanie modelu

# In[ ]:


def generate_model():
    model = Sequential()  # Linear stacking of layers

    # Convolution Layer 1
    model.add(Conv2D(16, (3, 3), input_shape=(32, 32, 3)))
    model.add(Activation("relu"))

    # ...

    model.add(Flatten())  # Flatten final output matrix into a vector

    # ...

    # Fully Connected Layer
    model.add(Dense(10))  # final 10 FC nodes
    model.add(Activation("softmax"))  # softmax activation

    model.summary()

    adam = tf.optimizers.Adam(learning_rate=0.001)
    model.compile(loss="categorical_crossentropy", optimizer=adam, metrics=["accuracy"])

    return model


def generate_model_default():
    model = Sequential()  # Linear stacking of layers

    # Convolution Layer 1
    model.add(Conv2D(16, (3, 3), input_shape=(32, 32, 3)))
    model.add(Activation("relu"))

    # ...

    model.add(Flatten())  # Flatten final output matrix into a vector

    # ...

    # Fully Connected Layer
    model.add(Dense(10))  # final 10 FC nodes
    model.add(Activation("softmax"))  # softmax activation

    model.summary()

    adam = tf.optimizers.Adam(learning_rate=0.001)
    model.compile(loss="categorical_crossentropy", optimizer=adam, metrics=["accuracy"])

    return model


def generate_model_gemini():
    model = Sequential()

    # Convolutional Layers with Max Pooling
    model.add(Conv2D(32, (3, 3), activation="relu", input_shape=(32, 32, 3)))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Conv2D(64, (3, 3), activation="relu"))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.25))  # Regularization

    # Flatten and Fully Connected Layers
    model.add(Flatten())
    model.add(Dense(128, activation="relu"))
    model.add(Dropout(0.5))  # Regularization
    model.add(Dense(10, activation="softmax"))

    # Model Compilation
    model.compile(
        loss="categorical_crossentropy",
        optimizer="adam",  # Consider trying other optimizers
        metrics=["accuracy"],
    )

    return model


def generate_model_chat():
    model = Sequential()  # Linear stacking of layers

    # Convolution Layer 1
    model.add(Conv2D(32, (3, 3), padding="same", input_shape=(32, 32, 3)))
    model.add(Activation("relu"))
    model.add(BatchNormalization())

    # Convolution Layer 2
    model.add(Conv2D(32, (3, 3)))
    model.add(Activation("relu"))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.3))

    # Convolution Layer 3
    model.add(Conv2D(64, (3, 3), padding="same"))
    model.add(Activation("relu"))
    model.add(BatchNormalization())

    # Convolution Layer 4
    model.add(Conv2D(64, (3, 3)))
    model.add(Activation("relu"))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.3))

    # Flattening the convolutions
    model.add(Flatten())

    # Fully Connected Layer
    model.add(Dense(512))  # Large fully connected layer
    model.add(Activation("relu"))
    model.add(BatchNormalization())
    model.add(Dropout(0.6))

    # Output Layer
    model.add(Dense(10))  # final 10 FC nodes
    model.add(Activation("softmax"))  # softmax activation

    model.summary()

    # Compile the model
    adam = Adam(learning_rate=0.001)
    model.compile(loss="categorical_crossentropy", optimizer=adam, metrics=["accuracy"])

    return model


# In[ ]:


model = generate_model()
model_default = generate_model_default()
model_gemini = generate_model_gemini()
model_chat = generate_model_chat()
models = [model_chat]


# ## Trening

# In[ ]:


gen = ImageDataGenerator(
    rotation_range=8,
    width_shift_range=0.08,
    shear_range=0.3,
    height_shift_range=0.08,
    zoom_range=0.08,
    validation_split=0.2,
)

train_generator = gen.flow(X_train, Y_train, batch_size=128, subset="training")
valid_generator = gen.flow(X_train, Y_train, batch_size=128, subset="validation")


# In[ ]:


# Max 20 epoch
for model in models:
    model.fit(
        train_generator,
        steps_per_epoch=40000 // 128,
        epochs=20,
        verbose=1,
        validation_data=valid_generator,
        validation_steps=10000 // 128,
    )


# ## Test

# In[ ]:

for model in models:
    score = model.evaluate(X_test, Y_test)
    print("Test score:", score[0])
    print("Test accuracy:", score[1])

    # The predict_classes function outputs the highest probability class
    # according to the trained classifier for each input example.
    predicted = model.predict(X_test)
    predicted_classes = np.argmax(predicted, axis=1)

    # Check which items we got right / wrong
    correct_indices = np.nonzero(predicted_classes == y_test)[0]

    incorrect_indices = np.nonzero(predicted_classes != y_test)[0]

    cnf_matrix = confusion_matrix(y_test, predicted_classes)

    class_names = [str(i) for i in range(10)]

    # Plot non-normalized confusion matrix
    plt.figure()
    plot_confusion_matrix(
        cnf_matrix, classes=class_names, title="Confusion matrix, without normalization"
    )

    plt.show()


# In[ ]:


def show_samples_rgb(indices, preds, images, labels, count=3, names=[]):
    plt.figure()
    for i, sample in enumerate(indices[: count**2]):
        pred_id = int(np.argmax(preds[sample]))
        real_id = int(labels[sample])
        pred_score = preds[sample][pred_id]
        real_score = preds[sample][real_id]
        plt.subplot(count, count, i + 1)
        plt.imshow(images[sample], interpolation="none")
        plt.axis("off")
        if len(names) > 0:
            plt.title(
                f"P: {names[pred_id]} ({pred_score:.2f})\nE: {names[real_id]} ({real_score:.2f})"
            )
        else:
            plt.title(
                f"P: {pred_id} ({pred_score:.2f})\nE: {real_id} ({real_score:.2f})"
            )

    plt.tight_layout()


# ## Poprawne klasyfikacje

# In[ ]:


show_samples_rgb(correct_indices, predicted, X_test, y_test, 5, cifar_names)


# ## Błędne klasyfikacje

# In[ ]:


show_samples_rgb(incorrect_indices, predicted, X_test, y_test, 5, cifar_names)