# TODO: Implement load the data here.
# Load pickled data
import pickle
import csv
import os

# TODO: fill this in based on where you saved the training and testing data
training_file = '../../traffic-signs/traffic-signs-data/train.p'
testing_file = '../../traffic-signs/traffic-signs-data/test.p'

with open(training_file, mode='rb') as f:
    train = pickle.load(f)
with open(testing_file, mode='rb') as f:
    test = pickle.load(f)

X_train, y_train = train['features'], train['labels']
X_test, y_test = test['features'], test['labels']

# Make dictionary of sign names from CSV file
with open('../../traffic-signs/signnames.csv', 'r') as csvfile:
    reader = csv.reader(csvfile)
    next(reader, None)  # skip the headers
    sign_names = dict((int(n),label) for n, label in reader)

cls_numbers, cls_names = zip(*sign_names.items())

n_classes = len(set(y_train))
flat_img_size = 32*32*3

# STOP: Do not change the tests below. Your implementation should pass these tests. 
assert(X_train.shape[0] == y_train.shape[0]), "The number of images is not equal to the number of labels."
assert(X_train.shape[1:] == (32,32,3)), "The dimensions of the images are not 32 x 32 x 3."

import numpy as np
from PIL import Image
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelBinarizer
from sklearn.utils import resample
from tqdm import tqdm_notebook
from zipfile import ZipFile
import time
from datetime import timedelta
import math
import tensorflow as tf

"""
Explanation: Traffic Sign Classification with Keras
Keras exists to make coding deep neural networks simpler. To demonstrate just how easy it is, you’re going to use Keras to build a convolutional neural network in a few dozen lines of code.
You’ll be connecting the concepts from the previous lessons to the methods that Keras provides.
Dataset
The network you'll build with Keras is similar to the example that you can find in Keras’s GitHub repository that builds out a convolutional neural network for MNIST. 
However, instead of using the MNIST dataset, you're going to use the German Traffic Sign Recognition Benchmark dataset that you've used previously.
You can download pickle files with sanitized traffic sign data here.
Overview
Here are the steps you'll take to build the network:

First load the data.
Build a feedforward neural network to classify traffic signs.
Build a convolutional neural network to classify traffic signs.

Keep an eye on the network’s accuracy over time. Once the accuracy reaches the 98% range, you can be confident that you’ve built and trained an effective model.
Load the Data
Start by importing the data from the pickle file.
End of explanation
"""


# TODO: Implement data normalization here.
def normalize_color(image_data):
    """
    Normalize the image data with Min-Max scaling to a range of [0.1, 0.9]
    :param image_data: The image data to be normalized
    :return: Normalized image data
    """
    a = -0.5
    b = +0.5
    
    Xmin = 0.0
    Xmax = 255.0

    norm_img = np.empty_like(image_data, dtype=np.float32)

    norm_img = a + (image_data - Xmin)*(b-a)/(Xmax - Xmin)
    
    return norm_img

X_train = normalize_color(X_train)
X_test = normalize_color(X_test)

# STOP: Do not change the tests below. Your implementation should pass these tests. 
assert(round(np.mean(X_train)) == 0), "The mean of the input data is: %f" % np.mean(X_train)
assert(np.min(X_train) == -0.5 and np.max(X_train) == 0.5), "The range of the input data is: %.1f to %.1f" % (np.min(X_train), np.max(X_train))

"""
Explanation: Normalize the data
Now that you've loaded the training data, normalize the input so that it has a mean of 0 and a range between -0.5 and 0.5.
End of explanation
"""


from keras.models import Sequential
from keras.layers.core import Dense, Activation
from keras.optimizers import Adam
from keras.utils import np_utils

# TODO: Build a two-layer feedforward neural network with Keras here.
model = Sequential()
model.add(Dense(128, input_shape=(flat_img_size,), name='hidden1'))
model.add(Activation('relu'))
model.add(Dense(43, name='output'))
model.add(Activation('softmax'))

# STOP: Do not change the tests below. Your implementation should pass these tests.
assert(model.get_layer(name="hidden1").input_shape == (None, 32*32*3)), "The input shape is: %s" % model.get_layer(name="hidden1").input_shape
assert(model.get_layer(name="output").output_shape == (None, 43)), "The output shape is: %s" % model.get_layer(name="output").output_shape 

"""
Explanation: Build a Two-Layer Feedfoward Network
The code you've written so far is for data processing, not specific to Keras. Here you're going to build Keras-specific code.
Build a two-layer feedforward neural network, with 128 neurons in the fully-connected hidden layer. 
To get started, review the Keras documentation about models and layers.
The Keras example of a Multi-Layer Perceptron network is similar to what you need to do here. Use that as a guide, but keep in mind that there are a number of differences.
End of explanation
"""


# One-Hot encode the labels
Y_train = np_utils.to_categorical(y_train, n_classes)
Y_test = np_utils.to_categorical(y_test, n_classes)

# Reshape input for MLP
X_train_mlp = X_train.reshape(-1, flat_img_size)
X_test_mlp = X_test.reshape(-1, flat_img_size)

model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

history = model.fit(X_train_mlp, Y_train, batch_size=128, nb_epoch=10,
                    validation_data=(X_test_mlp, Y_test), verbose=1)

# STOP: Do not change the tests below. Your implementation should pass these tests.
assert(history.history['acc'][0] > 0.5), "The training accuracy was: %.3f" % history.history['acc'][0]

"""
Explanation: Train the Network
Compile and train the network for 2 epochs. Use the adam optimizer, with categorical_crossentropy loss.
Hint 1: In order to use categorical cross entropy, you will need to one-hot encode the labels.
Hint 2: In order to pass the input images to the fully-connected hidden layer, you will need to reshape the input.
Hint 3: Keras's .fit() method returns a History.history object, which the tests below use. Save that to a variable named history.
End of explanation
"""


# Get randomized datasets for training and validation
X_train, X_val, Y_train, Y_val = train_test_split(
    X_train,
    Y_train,
    test_size=0.25,
    random_state=0xdeadbeef)

X_val_mlp = X_val.reshape(-1, flat_img_size)
print('Training features and labels randomized and split.')

# STOP: Do not change the tests below. Your implementation should pass these tests.
assert(round(X_train.shape[0] / float(X_val.shape[0])) == 3), "The training set is %.3f times larger than the validation set." % (X_train.shape[0] / float(X_val.shape[0]))
assert(history.history['val_acc'][0] > 0.6), "The validation accuracy is: %.3f" % history.history['val_acc'][0]

loss, acc = model.evaluate(X_val.reshape(-1, flat_img_size), Y_val, verbose=1)
print('\nValidation accuracy : {0:>6.2%}'.format(acc))

"""
Explanation: Validate the Network
Split the training data into a training and validation set.
Measure the validation accuracy of the network after two training epochs.
Hint: Use the train_test_split() method from scikit-learn.
End of explanation
"""


from keras.layers import Convolution2D, MaxPooling2D, Dropout, Flatten

# TODO: Re-construct the network and add a convolutional layer before the first fully-connected layer.
model = Sequential()
model.add(Convolution2D(16, 5, 5, border_mode='same', input_shape=(32, 32, 3)))
model.add(Activation('relu'))
model.add(Flatten())
model.add(Dense(128, input_shape=(flat_img_size,), name='hidden1'))
model.add(Activation('relu'))
model.add(Dense(43, name='output'))
model.add(Activation('softmax'))

# TODO: Compile and train the model.
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

history = model.fit(X_train, Y_train, batch_size=128, nb_epoch=10,
                    validation_data=(X_val, Y_val), verbose=1)

# STOP: Do not change the tests below. Your implementation should pass these tests.
assert(history.history['val_acc'][0] > 0.9), "The validation accuracy is: %.3f" % history.history['val_acc'][0]

"""
Explanation: Validation Accuracy: 95.92%
Congratulations
You've built a feedforward neural network in Keras!
Don't stop here! Next, you'll add a convolutional layer to drive.py.
Convolutions
Build a new network, similar to your existing network. Before the hidden layer, add a 3x3 convolutional layer with 32 filters and valid padding.
Then compile and train the network.
Hint 1: The Keras example of a convolutional neural network for MNIST would be a good example to review.
Hint 2: Now that the first layer of the network is a convolutional layer, you no longer need to reshape the input images before passing them to the network. You might need to reload your training data to recover the original shape.
Hint 3: Add a Flatten() layer between the convolutional layer and the fully-connected hidden layer.
End of explanation
"""


# TODO: Re-construct the network and add a pooling layer after the convolutional layer.
model = Sequential()
model.add(Convolution2D(16, 5, 5, border_mode='same', input_shape=(32, 32, 3)))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=(2,2)))
model.add(Flatten())
model.add(Dense(128, input_shape=(flat_img_size,), name='hidden1'))
model.add(Activation('relu'))
model.add(Dense(43, name='output'))
model.add(Activation('softmax'))

# TODO: Compile and train the model.
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

history = model.fit(X_train, Y_train, batch_size=128, nb_epoch=10,
                    validation_data=(X_val, Y_val), verbose=1)

# STOP: Do not change the tests below. Your implementation should pass these tests.
## Fixed bug
assert(history.history['val_acc'][-1] > 0.9), "The validation accuracy is: %.3f" % history.history['val_acc'][0]

"""
Explanation: Validation Accuracy: 96.98%
Pooling
Re-construct your network and add a 2x2 pooling layer immediately following your convolutional layer.
Then compile and train the network.
End of explanation
"""


# TODO: Re-construct the network and add dropout after the pooling layer.
model = Sequential()
model.add(Convolution2D(16, 5, 5, border_mode='same', input_shape=(32, 32, 3)))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=(2,2)))
model.add(Dropout(0.5))
model.add(Flatten())
model.add(Dense(128, input_shape=(flat_img_size,), name='hidden1'))
model.add(Activation('relu'))
model.add(Dense(43, name='output'))
model.add(Activation('softmax'))

# TODO: Compile and train the model.
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

history = model.fit(X_train, Y_train, batch_size=128, nb_epoch=10,
                    validation_data=(X_val, Y_val), verbose=1)

# STOP: Do not change the tests below. Your implementation should pass these tests.
assert(history.history['val_acc'][-1] > 0.9), "The validation accuracy is: %.3f" % history.history['val_acc'][0]

"""
Explanation: Validation Accuracy: 97.36%
Dropout
Re-construct your network and add dropout after the pooling layer. Set the dropout rate to 50%.
End of explanation
"""


pool_size = (2,2)
model = Sequential()
model.add(Convolution2D(16, 5, 5, border_mode='same', input_shape=(32, 32, 3)))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=pool_size))
model.add(Convolution2D(64, 3, 3))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=pool_size))
model.add(Dropout(0.5))
model.add(Convolution2D(128, 3, 3))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=pool_size))
model.add(Dropout(0.5))
model.add(Flatten())
model.add(Dense(256, input_shape=(flat_img_size,), name='hidden1'))
model.add(Activation('relu'))
model.add(Dense(43, name='output'))
model.add(Activation('softmax'))

model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

history = model.fit(X_train, Y_train, batch_size=128, nb_epoch=50,
                    validation_data=(X_val, Y_val), verbose=1)

"""
Explanation: Validation Accuracy: 97.75%
Optimization
Congratulations! You've built a neural network with convolutions, pooling, dropout, and fully-connected layers, all in just a few lines of code.
Have fun with the model and see how well you can do! Add more layers, or regularization, or different padding, or batches, or more training epochs.
What is the best validation accuracy you can achieve?
End of explanation
"""


# with open('./test.p', mode='rb') as f:
#     test = pickle.load(f)
    
# X_test = test['features']
# y_test = test['labels']
# X_test = X_test.astype('float32')
# X_test /= 255
# X_test -= 0.5
# Y_test = np_utils.to_categorical(y_test, 43)

model.evaluate(X_test, Y_test)

model.save('test-acc-9716-epoch50.h5')

from keras.models import load_model
model2 = load_model('test-acc-9716-epoch50.h5')
# model2.evaluate(X_test, Y_test)
model2.summary()

"""
Explanation: Best Validation Accuracy: 99.65%
Testing
Once you've picked out your best model, it's time to test it.
Load up the test data and use the evaluate() method to see how well it does.
Hint 1: After you load your test data, don't forget to normalize the input and one-hot encode the output, so it matches the training data.
Hint 2: The evaluate() method should return an array of numbers. Use the metrics_names() method to get the labels.
End of explanation
"""


