#!/usr/bin/python
# -*- coding: utf-8 -*-
import numpy as np
import cv2
import matplotlib.pyplot as plt
from scipy.misc import imread, imsave

def fit_from_sla(binary, lane_fit):

    nonzerox = np.array(binary.nonzero()[1])
    nonzeroy = np.array(binary.nonzero()[0])
    margin = 300
    lane_coor = (
        (nonzerox > (lane_fit[0] * (nonzeroy ** 2) + lane_fit[1] * nonzeroy + lane_fit[2] - margin)) & 
        (nonzerox < (lane_fit[0] * (nonzeroy ** 2) + lane_fit[1] * nonzeroy + lane_fit[2] + margin)))
    x = nonzerox[lane_coor]
    y = nonzeroy[lane_coor]
    out_img = binary.copy()
    out_img[nonzeroy[lane_coor], nonzerox[lane_coor]] = [190, 190, 255]

    try:
        lane_fit = np.polyfit(y, x, 2)
        ploty = np.linspace(0, binary.shape[0]-1, binary.shape[0])
        lane_fitx = lane_fit[0] * ploty ** 2 + lane_fit[1] * ploty + lane_fit[2]
        pts = np.array([np.transpose(np.vstack([lane_fitx, ploty]))])
        cv2.polylines(out_img, np.int32([pts]), False, (100, 100, 255), 20)
    except:
        lane_fit = None

    return out_img

def sliding_windows(binary, numwindows):

    nonzerox = np.array(binary.nonzero()[1])
    nonzeroy = np.array(binary.nonzero()[0])

    hist = np.sum(binary[np.int(binary.shape[0] / 2):, :], axis = 0)
    plt.plot(hist)
    plt.title('Histogram', fontsize=50, pad=30)
    plt.xticks(fontsize=20)
    plt.yticks(fontsize=20)
    plt.xlabel("Position [px]", fontsize=30, labelpad=10)
    plt.ylabel("Amount of white pixels [-]", fontsize=30, labelpad=10)
    plt.show()
    mid = np.int(hist.shape[0] / 2)
    
    x_base = np.mean(np.argmax(hist, axis=0), dtype=np.int32)
    #x_base = binary.shape[1] // 2
    x_current = x_base
    nwindows = 10
    margin = 300
    min_pix = 200
    window_height = np.int(binary.shape[0] / nwindows)
    lane_coor = []
    out_img = binary.copy()

    for w in range(numwindows):
        wy_low = binary.shape[0] - (w + 1) * window_height
        wy_upp = binary.shape[0] - w * window_height
        wx_low = x_current - margin
        wx_upp = x_current + margin
        cv2.rectangle(out_img, (wx_low, wy_low), (wx_upp, wy_upp), (0, 255, 0), 3)

        found_coor = (
            (nonzeroy >= wy_low) & (nonzeroy < wy_upp) & 
            (nonzerox >= wx_low) & (nonzerox < wx_upp)
            ).nonzero()[0]
        lane_coor.append(found_coor)

        if len(found_coor) > min_pix:
            x_current = np.int(np.mean(nonzerox[found_coor]))

    lane_coor = np.concatenate(lane_coor)
    out_img[nonzeroy[lane_coor], nonzerox[lane_coor]] = [190, 255, 190]
    x = nonzerox[lane_coor]
    y = nonzeroy[lane_coor]

    try:
        lane_fit = np.polyfit(y, x, 2)
        ploty = np.linspace(0, binary.shape[0]-1, binary.shape[0])
        lane_fitx = lane_fit[0] * ploty ** 2 + lane_fit[1] * ploty + lane_fit[2]
        pts = np.array([np.transpose(np.vstack([lane_fitx, ploty]))])
        cv2.polylines(out_img, np.int32([pts]), False, (0, 255, 0), 20)
    except:
        lane_fit = None

    return out_img, lane_fit

def show_images(img_list, labels):
    rows = 2
    cols = 3
    fig_size = (200, 100)
    figure, axes = plt.subplots(nrows=rows, ncols=cols, figsize=fig_size)
    counter = 0
    for i in range(0, rows):
        for j in range(0, cols):
            axes[i, j].imshow(img_list[counter], cmap='gray')
            #axes[i, j].set_title(labels[counter], fontsize=30, pad=15)
            axes[i, j].axis('off')
            counter +=1

    plt.tight_layout()
    plt.show()

nimgs = 6
img_list = [imread('../output_images/warped_image_alab_bin0' + str(i) + '.png') for i in range(1, nimgs + 1)]
label_list = ['warped_image_alab_bin0' + str(j) for j in range(1, nimgs + 1)]
out_img1, lane_fit1 = sliding_windows(img_list[0], 1)
out_img2, lane_fit2 = sliding_windows(img_list[0], 2)
out_img3, lane_fit3 = sliding_windows(img_list[0], 3)
out_img4, lane_fit4 = sliding_windows(img_list[0], 5)
out_img5, lane_fit5 = sliding_windows(img_list[0], 8)
out_img6, lane_fit6 = sliding_windows(img_list[0], 10)
"""
lane_fit = [lane_fit1, lane_fit2, lane_fit3, lane_fit4, lane_fit5, lane_fit6]
out_image = []
for i in range(len(img_list)):
    out_image.append(fit_from_sla(img_list[0], lane_fit[i]))
"""
out_img = [out_img1, out_img2, out_img3, out_img4, out_img5, out_img6]
#cv2.polylines(out_img[i])
show_images(out_img, label_list)




