from detection.lesson_functions import *
from lane_lines.nb_start import *

import cv2
# Define a function to extract features from a single image window
# This function is very similar to extract_features()
# just for a single image rather than list of images
def single_img_features(img, color_space='RGB', spatial_size=(32, 32),
                        hist_bins=32, orient=9,
                        pix_per_cell=8, cell_per_block=2, hog_channel=0):

    feature_image = convert_color(img, color_space)
    feature_image = cv2.resize(feature_image, (64, 64))
    # 3) Compute spatial features if flag is set

    spatial_features = bin_spatial(feature_image, size=spatial_size)
    hist_features = color_hist(feature_image, nbins=hist_bins)

    if hog_channel == 'ALL':
        hog_features = []
        for channel in range(feature_image.shape[2]):
            hog_features.extend(get_hog_features(feature_image[:, :, channel],
                                                 orient, pix_per_cell, cell_per_block,
                                                 vis=False, feature_vec=True))
    else:
        hog_features = get_hog_features(feature_image[:, :, hog_channel], orient,
                                        pix_per_cell, cell_per_block, vis=False,
                                        feature_vec=True)

    # 9) Return concatenated array of features
    return np.concatenate([spatial_features, hist_features, hog_features])


# Define a function you will pass an image
# and the list of windows to be searched (output of slide_windows())
def search_windows(img, windows, clf, color_space='RGB',
                   spatial_size=(32, 32), hist_bins=32, orient=9,
                   pix_per_cell=8, cell_per_block=2,
                   hog_channel=0, feature_lst = []):
    # 1) Create an empty list to receive positive detection windows
    on_windows = []
    feature_lst = []
    img_list = []
    # 2) Iterate over all windows in the list
    for window in windows:
        # 3) Extract the test window from original image
        test_img = cv2.resize(img[window[0][1]:window[1][1], window[0][0]:window[1][0]],
                              (64, 64))
        # 4) Extract features for that window using single_img_features()
        features = single_img_features(test_img, color_space=color_space,
                                       spatial_size=spatial_size, hist_bins=hist_bins,
                                       orient=orient, pix_per_cell=pix_per_cell,
                                       cell_per_block=cell_per_block,
                                       hog_channel=hog_channel)
        #feature_lst.append(features)
        # 6) Predict using your classifier

        prediction = clf.predict_proba(features.reshape(1, -1))[0][1]
        # 7) If positive (prediction == 1) then save the window
        #print(prediction)
        if prediction >= .5:
            feature_lst.append(features)
            img_list.append(test_img)

            # print(prediction)
            on_windows.append(window)
    # 8) Return windows for positive detections
    return on_windows, feature_lst, img_list
