#%%

import numpy as np
import matplotlib.pyplot as plt

def show_image(img):
    mn = np.min(img)
    mx = np.max(img)
    img2 = (img-mn)/(mx-mn)
    plt.imshow(img2, cmap="gray")

    
#%%

(training_set, test_set) = keras.datasets.cifar10.load_data()

#%%

(training_inputs, training_labels) = training_set
(test_inputs, test_labels) = test_set

# Scale images to the [0, 1] range
training_inputs = training_inputs.astype("float32") / 255
test_inputs = test_inputs.astype("float32") / 255

np.savez("cifar10", training_inputs=training_inputs, 
         training_labels=training_labels, 
         test_inputs=test_inputs, test_labels=test_labels)

#%%

cifar_data = np.load("cifar10.npz")
training_inputs = cifar_data["training_inputs"]
training_labels = cifar_data["training_labels"]
test_inputs = cifar_data["test_inputs"]
test_labels = cifar_data["test_labels"]

#%%

num_classes = training_labels.max()+1

class_indices = [None] * num_classes
for c in range(0, num_classes):
    (rows, cols) = np.nonzero(training_labels == c)
    class_indices[c] = rows
    
samples_per_class = 100
small_indices = class_indices[0][0:samples_per_class]
for c in range(1, num_classes):
    small_indices = np.concatenate((small_indices, class_indices[c][0:samples_per_class]))
    
small_training_inputs = training_inputs[small_indices]
small_training_labels = training_labels[small_indices]


input_shape = small_training_inputs[0].shape
model = keras.Sequential([
        keras.Input(shape=input_shape),
        layers.Conv2D(32, kernel_size=(3, 3), activation="sigmoid"),
        layers.Conv2D(32, kernel_size=(3, 3), activation="sigmoid"),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Conv2D(64, kernel_size=(3, 3), activation="sigmoid"),
        layers.Conv2D(64, kernel_size=(3, 3), activation="sigmoid"),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Flatten(),
        layers.Dropout(0.5),
        layers.Dense(num_classes, activation="softmax"),
    ])

model.summary()

batch_size = 32
epochs = 100

model.compile(loss=keras.losses.SparseCategoricalCrossentropy(), optimizer="adam", 
              metrics=["accuracy"])

callbacks = [
    keras.callbacks.ModelCheckpoint(
        filepath="cifar10_t100_b32.keras",
        save_best_only=True,
        monitor="val_loss")
    ]

# train the model on the "big classes" dataset
hist_tr1000_4conv_b32_sigmoid = model.fit(small_training_inputs, small_training_labels, 
                    epochs=epochs, batch_size=batch_size,
                    validation_data=[test_inputs, test_labels]
                    )

# test the model on the "big classes" test set.
test_loss, test_acc = model.evaluate(test_inputs, test_labels, verbose=0)
print('\nTest accuracy: %.2f' % (test_acc * 100))
#model.save('cifar10_tanh_t10_b4_e40.h5')

#%%

vgg16 = keras.applications.vgg16.VGG16(
    weights="imagenet",
    include_top=False,
    input_shape=input_shape)

# Here we do transfer learning.
#   This is the implementation that uses the Sequential API
#   - load the model that was trained on the "big" dataset
#   - remove the previous output layer
#   - add a new output layer
#   - freeze all weights except the ones in the new output layer
#   - train on small dataset

# load the model that was trained on the "big" dataset

num_layers = len(vgg16.layers)

# This is where we use the Sequential API to create the model.
# Notice that we create a list of layers, where we use all layers of the 
# model except for the last one, and we add a new fully connected
# output layer.
refined_model = keras.Sequential([keras.Input(shape=input_shape)]+
                                 vgg16.layers + 
                                 [layers.Flatten(),
                                  layers.Dropout(0.5),
                                  layers.Dense(512, activation="tanh"),
                                  layers.Dropout(0.5),
                                  layers.Dense(num_classes, activation="softmax")])

# We freeze the weights of the VGG-16 layers we are using.
num_base_layers = len(vgg16.layers)
for i in range(0, num_base_layers):
    refined_model.layers[i].trainable = False


refined_model.compile(loss=keras.losses.SparseCategoricalCrossentropy(), 
                      optimizer=keras.optimizers.Adam(), 
                      metrics=["accuracy"])
refined_model.summary()

# train the model (essentially, train the weights of the new layers we added) 
# on the "small" dataset. 
epochs = 100
batch_size = 32

hist_tr100_den2_512 = refined_model.fit(small_training_inputs, 
                                        small_training_labels, 
                                        epochs=epochs, 
                                        batch_size=batch_size,
                                        validation_data=(test_inputs, test_labels))

test_loss, test_acc = refined_model.evaluate(test_inputs, test_labels, verbose=0)
print('\nTest accuracy: %.2f' % (test_acc * 100))
#refined_model.save('cifar10_vgg16_tr10_b4_e20.h5')

#%%