import csv
import cv2
import numpy as np
import datetime
import random
from sklearn.utils import shuffle
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

lines = []

#load csv file
with open('mydata/driving_log.csv') as csvfile:
    reader = csv.reader(csvfile)
    header = next(reader)
    for line in reader:
        lines.append(line)

#crop resize and change color space of image
def change_color_space(image):
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return image

#generator to yeild processed images for training as well as validation data set
def generator_images(data, batchSize = 64):
    while True:
        data = shuffle(data)
        for i in range(0, len(data), int(batchSize/4)):
            X_batch = []
            y_batch = []
            details = data[i: i+int(batchSize/4)]
            for line in details:
                image = change_color_space(cv2.imread('mydata/IMG/' + line[0].split('/')[-1]))
                steering_angle = float(line[3])
                
                #appending original image
                X_batch.append(image)
                y_batch.append(steering_angle)
                
                #appending flipped image
                X_batch.append(np.fliplr(image))
                y_batch.append(-steering_angle)
                
                # appending left camera image and steering angle with offset
                X_batch.append(change_color_space(cv2.imread('mydata/IMG/' + line[1].split('/')[-1])))
                y_batch.append(steering_angle+0.4)
                
                # appending right camera image and steering angle with offset
                X_batch.append(change_color_space(cv2.imread('mydata/IMG/' + line[2].split('/')[-1])))
                y_batch.append(steering_angle-0.3)
            # converting to numpy array
            X_batch = np.array(X_batch)
            y_batch = np.array(y_batch)
            yield shuffle(X_batch, y_batch)

#Diving data among training and validation set
training_data, validatio_data = train_test_split(lines, test_size = 0.2)

from keras.models import Sequential
from keras.layers import Convolution2D, Dropout, MaxPooling2D, Flatten, Activation, Dense, Cropping2D, Lambda

#creating model to be trained
model = Sequential()
model.add(Lambda(lambda x: ((x/255.0) - 0.5), input_shape=(160,320,3) ))
model.add(Cropping2D(cropping=((70,25),(0,0))))
model.add(Convolution2D(24,5,5,subsample=(2,2),activation="relu"))
model.add(Convolution2D(36,5,5,subsample=(2,2),activation="relu"))
model.add(Convolution2D(48,5,5,subsample=(2,2),activation="relu"))
model.add(Convolution2D(64,3,3,activation="relu"))
model.add(Convolution2D(64,3,3,activation="relu"))
model.add(Dropout(0.5))
model.add(Flatten())
model.add(Dense(100))
model.add(Dense(50))
model.add(Dense(10))
model.add(Dense(1))

#compiling and running the model
model.compile('adam', 'mse')

history_object = model.fit_generator(generator_images(training_data), samples_per_epoch = len(training_data)*2, nb_epoch = 3, validation_data=generator_images(validatio_data), nb_val_samples=len(validatio_data), verbose=1)

### print the keys contained in the history object
print(history_object.history.keys())
print(history_object.history)

### plot the training and validation loss for each epoch
plt.plot(history_object.history['loss'])
plt.plot(history_object.history['val_loss'])
plt.title('model mean squared error loss')
plt.ylabel('mean squared error loss')
plt.xlabel('epoch')
plt.legend(['training set', 'validation set'], loc='upper right')
plt.show()

#saving the model
model.save('model.h5')