# Load pickled data
import pickle
import glob
import cv2
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import scipy as sp
import tensorflow as tf
import matplotlib.image as mpimg
from sklearn.model_selection import train_test_split
from sklearn.utils import shuffle
from tensorflow.contrib.layers import flatten
%matplotlib inline

training_file = 'data/train.p'
validation_file = 'data/valid.p'
testing_file = 'data/test.p'

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

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

### Replace each question mark with the appropriate value.
### Use python, pandas or numpy methods rather than hard coding the results

# TODO: Number of training examples
n_train = len(X_train)

# TODO: Number of validation examples
n_validation = len(X_valid)

# TODO: Number of testing examples.
n_test = len(X_test)

# TODO: What's the shape of an traffic sign image?
image_shape = X_train[0].shape

# TODO: How many unique classes/labels there are in the dataset.
n_classes = len(set(y_train))

print("Number of training examples =", n_train)
print("Number of testing examples =", n_test)
print("Image data shape =", image_shape)
print("Number of classes =", n_classes)

sign_names = pd.read_csv('signnames.csv', index_col='ClassId')
sign_names

for i in sorted(set(y_train)):
    print("Class ID: {}, Name: {}".format(i, sign_names.iloc[i]['SignName']))
    X_train_, y_train_ = shuffle(X_train, y_train)
    f, ax = plt.subplots(1, 5, figsize=(20,10))
    for image, i in zip(X_train_[y_train_ == i], range(5)):
        ax[i].imshow(image)
        ax[i].set_axis_off()
    plt.show()

from collections import Counter
label_count = Counter(y_train)
uniq_labels = sorted(set(y_train))
label_height = []
for i in uniq_labels:
    label_height.append(label_count[i])
plt.bar(uniq_labels, label_height)
plt.show()
print(label_count)

def convert_gray(X):
    X_ = []
    for x in X:
        X_.append(cv2.cvtColor(x, cv2.COLOR_RGB2GRAY))
    return np.array(X_)

def blur(X, y, kernel_size=5):
    X_ = []
    y_ = []
    for i in range(len(X)):
        X_.append(cv2.GaussianBlur(X[i], (kernel_size, kernel_size), 0))
        y_.append(y[i])
    return X_, y_

def rotate(X, y, angle=10):
    X_ = []
    y_ = []
    for i in range(len(X)):
        X_.append(sp.ndimage.interpolation.rotate(X[i], angle, reshape=False, mode='nearest'))
        y_.append(y[i])
    return X_, y_

def noise(X, y, mu=0, s=0.1):
    X_ = []
    y_ = []
    img_size = X.shape
    for i in range(len(X)):
        n = np.random.normal(mu, s, img_size[1]*img_size[2]).reshape(img_size[1:])
        X_.append(X[i]+n)
        y_.append(y[i])
    return X_, y_

def flip_x(X, y):
    X_ = []
    y_ = []
    for i in range(len(X)):
        X_.append(np.fliplr(X[i]))
        y_.append(y[i])
    return X_, y_

def normalize(x):
    return (np.float32(x)-128)/128

def extra_examples(X_train, y_train):
    X_extra = []
    y_extra = []

    X_extra_tmp, y_extra_tmp = rotate(X_train, y_train, 10)
    X_extra += X_extra_tmp
    y_extra += y_extra_tmp

    X_extra_tmp, y_extra_tmp = rotate(X_train, y_train, -10)
    X_extra += X_extra_tmp
    y_extra += y_extra_tmp

    X_extra_tmp, y_extra_tmp = rotate(X_train, y_train, 20)
    X_extra += X_extra_tmp
    y_extra += y_extra_tmp

    X_extra_tmp, y_extra_tmp = rotate(X_train, y_train, -20)
    X_extra += X_extra_tmp
    y_extra += y_extra_tmp

    X_extra_tmp, y_extra_tmp = noise(X_train, y_train)
    X_extra += X_extra_tmp
    y_extra += y_extra_tmp

    X_extra_tmp, y_extra_tmp = flip_x(X_train, y_train)
    X_extra += X_extra_tmp
    y_extra += y_extra_tmp

    X_extra_tmp, y_extra_tmp = blur(X_train, y_train)
    X_extra += X_extra_tmp
    y_extra += y_extra_tmp

    X_train = np.concatenate((X_train, X_extra), axis=0)
    y_train = np.concatenate((y_train, y_extra), axis=0)

    return X_train, y_train

def pipeline(X, y, extra=False):
    X = convert_gray(X)
    if extra:
        X, y = extra_examples(X, y)
    X = normalize(X)
    X = np.reshape(X, (X.shape[0], X.shape[1], X.shape[2], 1))
    return X, y

X_train, y_train = pipeline(X_train, y_train, extra=True)
X_valid, y_valid = pipeline(X_valid, y_valid)
X_test, y_test = pipeline(X_test, y_test)

X_train, y_train = shuffle(X_train, y_train)

EPOCHS = 40
BATCH_SIZE = 128

def LeNet(x):    
    # Arguments used for tf.truncated_normal, randomly defines variables for the weights and biases for each layer
    mu = 0
    sigma = 0.1

    # SOLUTION: Layer 1: Convolutional. Input = 32x32x1. Output = 28x28x6.
    conv1_W = tf.Variable(tf.truncated_normal(shape=(5, 5, 1, 6), mean = mu, stddev = sigma))
    conv1_b = tf.Variable(tf.zeros(6))
    conv1   = tf.nn.conv2d(x, conv1_W, strides=[1, 1, 1, 1], padding='VALID') + conv1_b

    # SOLUTION: Activation.
    conv1 = tf.nn.relu(conv1)

    # SOLUTION: Pooling. Input = 28x28x6. Output = 14x14x6.
    conv1 = tf.nn.max_pool(conv1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='VALID')

    # SOLUTION: Layer 2: Convolutional. Output = 10x10x16.
    conv2_W = tf.Variable(tf.truncated_normal(shape=(5, 5, 6, 16), mean = mu, stddev = sigma))
    conv2_b = tf.Variable(tf.zeros(16))
    conv2   = tf.nn.conv2d(conv1, conv2_W, strides=[1, 1, 1, 1], padding='VALID') + conv2_b

    # SOLUTION: Activation.
    conv2 = tf.nn.relu(conv2)

    # SOLUTION: Pooling. Input = 10x10x16. Output = 5x5x16.
    conv2 = tf.nn.max_pool(conv2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='VALID')

    # SOLUTION: Flatten. Input = 5x5x16. Output = 400.
    fc0   = flatten(conv2)

    # SOLUTION: Layer 3: Fully Connected. Input = 400. Output = 120.
    fc1_W = tf.Variable(tf.truncated_normal(shape=(400, 120), mean = mu, stddev = sigma))
    fc1_b = tf.Variable(tf.zeros(120))
    fc1   = tf.matmul(fc0, fc1_W) + fc1_b

    # SOLUTION: Activation.
    fc1    = tf.nn.relu(fc1)

    # DROPOUT    ####
    keep_prob = tf.placeholder(tf.float32)
    fc1_drop = tf.nn.dropout(fc1, keep_prob)

    # SOLUTION: Layer 4: Fully Connected. Input = 120. Output = 84.
    fc2_W  = tf.Variable(tf.truncated_normal(shape=(120, 84), mean = mu, stddev = sigma))
    fc2_b  = tf.Variable(tf.zeros(84))
    fc2    = tf.matmul(fc1_drop, fc2_W) + fc2_b   ####

    # SOLUTION: Activation.
    fc2    = tf.nn.relu(fc2)

    # SOLUTION: Layer 5: Fully Connected. Input = 84. Output = 10.
    fc3_W  = tf.Variable(tf.truncated_normal(shape=(84, 43), mean = mu, stddev = sigma))
    fc3_b  = tf.Variable(tf.zeros(43))
    logits = tf.matmul(fc2, fc3_W) + fc3_b

    return logits, keep_prob

x = tf.placeholder(tf.float32, (None, 32, 32, 1))
y = tf.placeholder(tf.int32, (None))
one_hot_y = tf.one_hot(y, 43)

rate = 0.001

logits, keep_prob = LeNet(x)
cross_entropy = tf.nn.softmax_cross_entropy_with_logits(labels=one_hot_y, logits=logits)
loss_operation = tf.reduce_mean(cross_entropy)
optimizer = tf.train.AdamOptimizer(learning_rate = rate)
training_operation = optimizer.minimize(loss_operation)

correct_prediction = tf.equal(tf.argmax(logits, 1), tf.argmax(one_hot_y, 1))
accuracy_operation = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
saver = tf.train.Saver()

def evaluate(X_data, y_data):
    num_examples = len(X_data)
    total_accuracy = 0
    sess = tf.get_default_session()
    for offset in range(0, num_examples, BATCH_SIZE):
        batch_x, batch_y = X_data[offset:offset+BATCH_SIZE], y_data[offset:offset+BATCH_SIZE]
        accuracy = sess.run(accuracy_operation, feed_dict={x: batch_x, y: batch_y, keep_prob: 1.0})
        total_accuracy += (accuracy * len(batch_x))
    return total_accuracy / num_examples

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    num_examples = len(X_train)

    print("Training...")
    print()
    for i in range(EPOCHS):
        X_train, y_train = shuffle(X_train, y_train)
        for offset in range(0, num_examples, BATCH_SIZE):
            end = offset + BATCH_SIZE
            batch_x, batch_y = X_train[offset:end], y_train[offset:end]
            sess.run(training_operation, feed_dict={x: batch_x, y: batch_y, keep_prob: 0.5})

        test_accuracy = evaluate(X_train, y_train)
        validation_accuracy = evaluate(X_valid, y_valid)
        print("EPOCH {} ...".format(i+1))
        print("Test Accuracy = {:.3f}".format(test_accuracy))
        print("Validation Accuracy = {:.3f}".format(validation_accuracy))
        print()

    saver.save(sess, './lenet')
    print("Model saved")

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    saver.restore(sess, './lenet')
    train_accuracy = evaluate(X_train, y_train)
    validation_accuracy = evaluate(X_valid, y_valid)
    test_accuracy = evaluate(X_test, y_test)
    print("Train Accuracy = {:.3f}".format(train_accuracy))
    print("Validation Accuracy = {:.3f}".format(validation_accuracy))
    print("Test Accuracy = {:.3f}".format(test_accuracy))
    print()

def resize(image):
    return cv2.resize(image, (32, 32))

wild = []
wild_filenames = [1, 13, 17, 25, 26, 38]
y_wild = wild_filenames
folder = 'wild'
for filename in wild_filenames:
    path = '{}/{}.jpg'.format(folder, filename)
    image = mpimg.imread(path)
    image_32 = resize(image)
    wild.append(image_32)
    f, ax = plt.subplots(1, 1, figsize=(2,2))
    ax.set_title(filename)
    ax.imshow(image_32)
    plt.show()

wild_gray = convert_gray(wild)
for image in wild_gray:
    plt.imshow(image, cmap='gray')
    plt.show()

X_wild = normalize(wild_gray)
X_wild = np.reshape(X_wild, (X_wild.shape[0], X_wild.shape[1], X_wild.shape[2], 1))

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    saver.restore(sess, './lenet')
    prediction = tf.argmax(tf.nn.softmax(logits),1)
    wild_predictions = sess.run(prediction, feed_dict={x: X_wild, keep_prob: 1.0})
    wild_softmax = sess.run(tf.nn.softmax(logits), feed_dict={x: X_wild, keep_prob: 1.0})
    print("Labels:     ", np.array(y_wild))
    print("Predictions:", wild_predictions)
    wild_accuracy = evaluate(X_wild, y_wild)
    print("Accuracy = {:.3f}".format(wild_accuracy))
    print()

### Calculate the accuracy for these 5 new images.
### For example, if the model predicted 1 out of 5 signs correctly, it's 20% accurate on these new images.
print(wild_predictions)
wild_accuracy_percent = np.sum(y_true == y_pred for y_true, y_pred in zip(y_wild, wild_predictions))/len(y_wild)*100
print("Accuracy: {}%".format(wild_accuracy_percent))

with tf.Session() as sess:
    res = sess.run(tf.nn.top_k(tf.constant(wild_softmax), k=5))
    print(res)