#Create references to important directories we will use over and over
import os, sys
DATA_HOME_DIR = '/home/nathan/olin/spring2017/line-follower/line-follower/data'

#import modules
import numpy as np
from glob import glob
from PIL import Image
from tqdm import tqdm
from scipy.ndimage import zoom

from keras.models import Sequential
from keras.metrics import categorical_crossentropy, categorical_accuracy
from keras.layers.convolutional import *
from keras.preprocessing import image
from keras.layers.core import Flatten, Dense, Reshape
from keras.optimizers import Adam
from keras.layers.recurrent import LSTM
from keras.layers.normalization import BatchNormalization

from matplotlib import pyplot as plt
import seaborn as sns
%matplotlib inline

"""
Explanation: Line Follower - CompRobo17
LSTM
Imports and Directories
End of explanation
"""


%cd $DATA_HOME_DIR

path = DATA_HOME_DIR
train_path=path + '/sun_apr_16_office_full_line_1'
valid_path=path + '/sun_apr_16_office_full_line_2'

"""
Explanation: Create paths to data directories
End of explanation
"""


def resize_vectorized4D(data, new_size=(64, 64)):
    """
    A vectorized implementation of 4d image resizing
    
    Args:
        data (4D array): The images you want to resize
        new_size (tuple): The desired image size
        
    Returns: (4D array): The resized images
    """
    fy, fx = np.asarray(new_size, np.float32) / data.shape[1:3]
    return zoom(data, (1, fy, fx, 1), order=1) # order is the order of spline interpolation

def lowerHalfImage(array):
    """ 
    Returns the lower half rows of an image
    
    Args: array (array): the array you want to extract the lower half from
    
    Returns: The lower half of the array
    """
    return array[round(array.shape[0]/2):,:,:]

def folder_to_numpy(image_directory_full):
    """
    Read sorted pictures (by filename) in a folder to a numpy array. 
    We have hardcoded the extraction of the lower half of the images as
    that is the relevant data
    
    USAGE:
        data_folder = '/train/test1'
        X_train = folder_to_numpy(data_folder)
    
    Args:
        data_folder (str): The relative folder from DATA_HOME_DIR
        
    Returns:
        picture_array (np array): The numpy array in tensorflow format
    """
    # change directory
    print ("Moving to directory: " + image_directory_full)
    os.chdir(image_directory_full)
    
    # read in filenames from directory
    g = glob('*.png')
    if len(g) == 0:
        g = glob('*.jpg')
    print ("Found {} pictures".format(len(g)))
    
    # sort filenames
    g.sort()
    
    # open and convert images to numpy array - then extract the lower half of each image
    print("Starting pictures to numpy conversion")
    picture_arrays = np.array([lowerHalfImage(np.array(Image.open(image_path))) for image_path in g])
    
#     reshape to tensorflow format
#     picture_arrays = picture_arrays.reshape(*picture_arrays.shape, 1)
    print ("Shape of output: {}".format(picture_arrays.shape))
    
    # return array
    return picture_arrays
    return picture_arrays.astype('float32')

def flip4DArray(array):
    """ Produces the mirror images of a 4D image array """
    return array[..., ::-1,:] #[:,:,::-1] also works but is 50% slower

def concatCmdVelFlip(array):
    """ Concatentaes and returns Cmd Vel array """
    return np.concatenate((array, array*-1)) # multiply by negative 1 for opposite turn

"""
Explanation: Helper Functions
Throughout the notebook, we will take advantage of helper functions to cleanly process our data.
End of explanation
"""


%cd $train_path
Y_train = np.genfromtxt('cmd_vel.csv', delimiter=',')[:,1] # only use turning angle
Y_train = np.concatenate((Y_train, Y_train*-1))
X_train = folder_to_numpy(train_path + '/raw')
X_train = np.concatenate((X_train, flip4DArray(X_train)))

"""
Explanation: Data
Because we are using a CNN and unordered pictures, we can flip our data and concatenate it on the end of all training and validation data to make sure we don't bias left or right turns.
Training Data
Extract and store the training data in X_train and Y_train
End of explanation
"""


X_train.shape, Y_train.shape

"""
Explanation: Test the shape of the arrays:
X_train: (N, 240, 640, 3)
Y_train: (N,)
End of explanation
"""


%cd /tmp
img = Image.fromarray(X_train[0], 'RGB')
img.save("temp.jpg")
image.load_img("temp.jpg")

"""
Explanation: Visualize the training data, currently using a hacky method to display the numpy matrix as this is being run over a remote server and I can't view new windows
End of explanation
"""


%cd $valid_path
Y_valid = np.genfromtxt('cmd_vel.csv', delimiter=',')[:,1]
Y_valid = np.concatenate((Y_valid, Y_valid*-1))
X_valid = folder_to_numpy(valid_path + '/raw')
X_valid = np.concatenate((X_valid, flip4DArray(X_valid)))

"""
Explanation: Validation Data
Follow the same steps for as the training data for the validation data.
End of explanation
"""


X_valid.shape, Y_valid.shape

"""
Explanation: Test the shape of the arrays:
X_valid: (N, 240, 640, 3)
Y_valid: (N,)
End of explanation
"""


img_rows, img_cols = (64, 64)

print(img_rows)
print(img_cols)

X_train = resize_vectorized4D(X_train, (img_rows, img_cols))
X_valid = resize_vectorized4D(X_valid, (img_rows, img_cols))

print(X_train.shape)
print(X_valid.shape)

"""
Explanation: Resize Data
When we train the network, we don't want to be dealing with (240, 640, 3) images as they are way too big. Instead, we will resize the images to something more managable, like (64, 64, 3) or (128, 128, 3). In terms of network predictive performance, we are not concerned with the change in aspect ratio, but might want to test a (24, 64, 3) images for faster training
End of explanation
"""


%cd /tmp
img = Image.fromarray(X_train[np.random.randint(0, X_train.shape[0])], 'RGB')
img.save("temp.jpg")
image.load_img("temp.jpg")

"""
Explanation: Visualize newly resized image.
End of explanation
"""


X_train = X_train / 256
X_valid = X_valid / 256

"""
Explanation: Rescale Images
End of explanation
"""


def windows(X, Y, seq_len=10):
    result = []
    for index in range(X.shape[0] - seq_len):
        result.append([X[index:index+seq_len], Y[index+seq_len-1]])
    result = np.array(result)
    np.random.shuffle(result)
    return np.array(list(result[:,0])), np.array(list(result[:,1]))

X_train_windows, Y_train_windows = windows(X_train, Y_train, 10)
X_valid_windows, Y_valid_windows = windows(X_valid, Y_valid, 10)

X_train_windows.shape

Y_train_windows.shape

# gen = image.ImageDataGenerator(
#                                 rescale=1. / 255 # normalize data between 0 and 1
#                               )

"""
Explanation: Batches
End of explanation
"""


# batch_size = 3
# train_generator = gen.flow(X_train_windows, Y_train_windows, batch_size=batch_size, shuffle=True)
# valid_generator = gen.flow(X_valid_windows, Y_valid_windows, batch_size=batch_size, shuffle=True)

# get_batches(train_path, batch_size=batch_size, 
#                             target_size=in_shape, 
#                             gen=gen)
# val_batches   = get_batches(valid_path, batch_size=batch_size, 
#                             target_size=in_shape, 
#                             gen=gen)

# data, category = next(train_generator)
# print ("Shape of data: {}".format(data[0].shape))
# %cd /tmp
# img = Image.fromarray(data[np.random.randint(0, data.shape[0])].astype('uint8'), 'RGB')
# img.save("temp.jpg")
# image.load_img("temp.jpg")

"""
Explanation: Next, create the train and valid generators, these are shuffle and have a batch size of 32 by default
End of explanation
"""


in_shape = (10, img_rows, img_cols, 3)

"""
Explanation: Convnet
Constants
End of explanation
"""


def get_model():
    model = Sequential([
        Conv3D(32, 1,3,3, border_mode='same', activation='relu', input_shape=in_shape),
        MaxPooling3D(pool_size=(1,2,2)),
        Conv3D(64, 1,3,3, border_mode='same', activation='relu'),
        MaxPooling3D(pool_size=(1,2,2)),
        Conv3D(128, 1,3,3, border_mode='same', activation='relu'),
        MaxPooling3D(pool_size=(1,2,2)),
#         Convolution2D(64,3,3, border_mode='same', activation='relu'),
#         MaxPooling2D(),
#         Convolution2D(128,3,3, border_mode='same', activation='relu'),
#         MaxPooling2D(),
#         Reshape((10, 8192)),
        Flatten(),
#         LSTM(100), #return_sequences=True
        Dense(2048, activation='relu'),
        Dense(1024, activation='relu'),
        Dense(512, activation='relu'),
        Dense(1)
        ])
    model.compile(loss='mean_absolute_error', optimizer='adam')
    return model

model = get_model()

model.summary()

"""
Explanation: Model
Our test model will use a VGG like structure with a few changes. We are removing the final activation function. We will also use either mean_absolute_error or mean_squared_error as our loss function for regression purposes.
End of explanation
"""


history = model.fit(X_train_windows, Y_train_windows,
                    batch_size = 3,
                    nb_epoch=8*60*2,
                    validation_data=(X_valid_windows, Y_valid_windows),
                    verbose=True)

# %cd $DATA_HOME_DIR
# model.save_weights('epoche_stateful_800.h5')

"""
Explanation: Train
End of explanation
"""


val_plot = np.convolve(history.history['val_loss'], np.repeat(1/10, 10), mode='valid')
train_plot = np.convolve(history.history['loss'], np.repeat(1/10, 10), mode='valid')

fig, ax = plt.subplots()
sns.tsplot(history.history['val_loss'], ax=ax)
ax.set_ylim([0,None])

X_preds = model.predict(X_valid_windows).reshape(X_valid_windows.shape[0],)
for i in range(len(X_valid)):
    print("{:07f} | {:07f}".format(Y_valid_windows[i], X_preds[i]))

X_train_preds = model.predict(X_train).reshape(X_train.shape[0],)
for i in range(len(X_train_preds)):
    print("{:07f} | {:07f}".format(Y_train[i], X_train_preds[i]))

"""
Explanation: Visualize Training
End of explanation
"""


