import numpy as np
import cv2
import scipy.misc
import pandas as pd
import matplotlib.pyplot as plt
from scipy.ndimage import rotate
from scipy.stats import bernoulli
import math

#useful constants
logFile3 = './data_u/driving_log3.csv'
logFile2 = './data_u/driving_log2.csv'
logFile1 = './data_u/driving_log1.csv'
imgPath = './data_u/'
steeringCoef = 0.229

def shear(image, steering_angle, shear_range=200):
     rows, cols, ch = image.shape
     dx = np.random.randint(-shear_range, shear_range + 1)
     random_point = [cols / 2 + dx, rows / 2]
     pts1 = np.float32([[0, rows], [cols, rows], [cols / 2, rows / 2]])
     pts2 = np.float32([[0, rows], [cols, rows], random_point])
     dsteering = dx / (rows / 2) * 360 / (2 * np.pi * 25.0) / 6.0
     M = cv2.getAffineTransform(pts1, pts2)
     image = cv2.warpAffine(image, M, (cols, rows), borderMode=1)
     steering_angle += dsteering
 
     return image, steering_angle


def crop(image, top, bottom):
    """
    top: portion of the image be cropped from the top of the image

    bottom: portion of the image cropped from the bottom of the image

    """
    assert 0 <= top < 0.5, 'should be between 0.0 and 0.5'
    assert 0 <= bottom < 0.5, 'should be between 0.0 and 0.5'

    top = int(np.ceil(image.shape[0] * top))
    bottom = image.shape[0] - int(np.ceil(image.shape[0] * bottom))

    return image[top:bottom, :]

def flip(image, steering_angle, flipping_prob=0.8):
    """
    Based on the outcome of an coin flip, the image will be flipped.
    If flipping is applied, the steering angle will be negated.
    :param steering_angle: Original steering angle
    :return: Both flipped image and new steering angle
    """
    head = bernoulli.rvs(flipping_prob)
    if head:
        return np.fliplr(image), -1 * steering_angle
    else:
        return image, steering_angle

def brightness(image):
    """
    Random gamma correction is used as an alternative method changing the brightness of
    training images.
    http://www.pyimagesearch.com/2015/10/05/opencv-gamma-correction/

    returns a new image generated by applying gamma correction to the source image
    """
    gamma = np.random.uniform(0.4, 1.5)
    inv_gamma = 1.0 / gamma
    table = np.array([((i / 255.0) ** inv_gamma) * 255
                      for i in np.arange(0, 256)]).astype("uint8")

    # apply gamma correction using the lookup table
    return cv2.LUT(image, table)


def generate_new_image(image, steering_angle, top_crop=0.35, bottom_crop=0.1,
                       resize_dim=(64, 64), do_shear_prob=0.9):
    
    head = bernoulli.rvs(do_shear_prob)
    if head == 1:
        image, steering_angle = shear(image, steering_angle)

    image = crop(image, top_crop, bottom_crop)

    image, steering_angle = flip(image, steering_angle)

    image = brightness(image)

    image = scipy.misc.imresize(image, resize_dim)

    return image, steering_angle

def get_next_image_files(batch_size=64):
    """
    The simulator records three images (namely: left, center, and right) at a given time
    However, when we are picking images for training we randomly (with equal probability)
    one of these three images and its steering angle.

    batch_size: Size of the image batch

    returns a list of selected (image files names, respective steering angles)
    """
    data1 = pd.read_csv(logFile1)
    data2 = pd.read_csv(logFile2)
    data3 = pd.read_csv(logFile3)
    data  = pd.concat([data1, data2, data3])
    data = pd.read_csv(logFile1)
    num_of_img = len(data)
    rnd_indices = np.random.randint(0, num_of_img, batch_size)

    image_files_and_angles = []
    for index in rnd_indices:
        if pd.isnull(data.iloc[index]['left']):
            rnd_image = 1
        else:
            rnd_image = np.random.randint(0, 3)
        
        if rnd_image == 0:
            img = data.iloc[index]['left'].strip()
            angle = data.iloc[index]['steering'] + steeringCoef
            image_files_and_angles.append((img, angle))

        elif rnd_image == 1:
            img = data.iloc[index]['center'].strip()
            angle = data.iloc[index]['steering']
            image_files_and_angles.append((img, angle))
        else:
            img = data.iloc[index]['right'].strip()
            angle = data.iloc[index]['steering'] - 0.2#steeringCoef
            image_files_and_angles.append((img, angle))

    return image_files_and_angles


def generate_next_batch(batch_size=64):
    """
    This generator yields the next training batch

    param batch_size: Number of training images in a single batch

    returns a tuple of features and steering angles as two numpy arrays
    """
    while True:
        X_batch = []
        y_batch = []
        images = get_next_image_files(batch_size)
        for img_file, angle in images:
            raw_image = plt.imread(imgPath + img_file)
            raw_angle = angle
            new_image, new_angle = generate_new_image(raw_image, raw_angle)
            X_batch.append(new_image)
            y_batch.append(new_angle)

        assert len(X_batch) == batch_size, 'len(X_batch) == batch_size should be True'

        yield np.array(X_batch), np.array(y_batch)

