import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

"""
Explanation: Goal
Create a network that detects cats.
End of explanation
"""


# Create image generators that yield data for creating the feature cache.
from keras.preprocessing.image import ImageDataGenerator
from good_dog.models import vgg16_image_preprocessor

train_generator = ImageDataGenerator(
    preprocessing_function=vgg16_image_preprocessor
)

test_generator = ImageDataGenerator(
    preprocessing_function=vgg16_image_preprocessor
)

# Define some settings for our model
from good_dog.models import Params

params = Params(batch_size=16, image_size=(150, 150))
classes = ["cats", "dogs"]
n_classes = len(classes)

# Create the bottom layers for VGG16, excluding any top dense layers.
from keras.applications import VGG16

vgg16_bottom = VGG16(include_top=False, weights='imagenet')

# Create or load our feature cache.
# Creating the feature cache can take a long time, so it's best to download when possible.
from good_dog.models import FeatureCache
from good_dog.data import require_directory

require_directory("cache/")
train_feature_cache = FeatureCache(params=params, 
                                   model=vgg16_bottom, 
                                   generator=train_generator, 
                                   directory="../data/images/train", 
                                   count=5000, 
                                   cache_path="cache/dog_v_cat_vgg16_train_features.pkl")
train_features = train_feature_cache.features
train_labels = train_feature_cache.labels

test_feature_cache = FeatureCache(params=params, 
                                   model=vgg16_bottom, 
                                   generator=test_generator, 
                                   directory="../data/images/validate", 
                                   count=800, 
                                   cache_path="cache/dog_v_cat_vgg16_test_features.pkl")
test_features = test_feature_cache.features
test_labels = test_feature_cache.labels

# Create the top model where we use the pretrained model from vgg16.
# This feeds into a GAP layer with some dropout to prevent overfitting.
from keras.layers import GlobalAveragePooling2D
from keras.layers import Dense
from keras.layers import Flatten
from keras.layers import Dropout
from keras.layers import BatchNormalization
from keras.models import Sequential

top_model = Sequential()
top_model.add(GlobalAveragePooling2D(input_shape=train_features.shape[1:]))
top_model.add(Dropout(0.5))
top_model.add(Dense(n_classes, activation='softmax'))

import keras.optimizers as optimizers
import keras.losses as losses

# Compile our model and set optimization strategy.
optimizer = optimizers.Adam()
top_model.compile(loss=losses.categorical_crossentropy, 
                  optimizer=optimizer, 
                  metrics=['accuracy'])
# Print out our model below.
top_model.summary()

# Train our model.
top_model.fit(train_features,
              train_labels,
              epochs=10,
              validation_data=(test_features, test_labels));

"""
Explanation: Load Data
Data comes from Dogs vs Cats dataset.

The way Keras looks for data is in a data/images/&lt;class-name&gt;/*.jpg path. The preprocessor in bin/dogs_vs_cats_preprocessor.py can create this format from the train file provided by the competition.
However the notebook here uses cached features if available, and it's ideal to use these as they significantly speed up training.
Define Network
To be sample efficient and generalize well, start with a VGG16 Network trained on the large ImageNet dataset. It has cats and dogs as part of its training distribution so it likely already has a good grasp of the features essential in detecting a cat.
End of explanation
"""


# Combine the bottom and top so we can predict all the way through.
from good_dog.models import MultiModel

model = MultiModel([vgg16_bottom, top_model])

# Show output prediction for an exemplar image
from keras.preprocessing.image import load_img

image_path = "../data/external/cat.jpg"
img = load_img(image_path, target_size=params.image_size)
plt.imshow(img)
vgg_img = test_generator.flow(np.expand_dims(img, axis=0)).next()
print("Score: %s" % list(zip(classes, model.predict(vgg_img).flat)))

"""
Explanation: Validate Model
Check our performance on cats vs dogs.
End of explanation
"""


# Show class activation map for our image.
from good_dog.models import ClassActivationMap

image_path = "../data/external/dog.jpg"
img = load_img(image_path, target_size=params.image_size)
plt.imshow(img)
vgg_img = test_generator.flow(np.expand_dims(img, axis=0)).next()
print("Score: %s" % list(zip(classes, model.predict(vgg_img).flat)))

gap_layer = top_model.layers[0]
dense_layer = top_model.layers[-1]
cam = ClassActivationMap(top_model.inputs, gap_layer, dense_layer)
vgg16_feature = model.models[0].predict(vgg_img)

for class_index, name in enumerate(classes):
    softmax_activation_map = cam.softmax_map_for_class(vgg16_feature, class_index)
    grayscale = np.array(img).mean(axis=2)
    plt.title("%s softmax activation map" % name)
    plt.imshow(grayscale, cmap='gray', vmin=0, vmax=255, alpha=0.5, origin='upper')
    extent = (0, params.image_size[0], params.image_size[0], 0)
    plt.imshow(softmax_activation_map, cmap='RdYlBu', vmax=1, vmin=0, extent=extent, alpha=0.5, origin='upper')
    plt.show()

"""
Explanation: Visualize Model
Show neurons that activate highly in GAP layer.
End of explanation
"""


