import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import numpy as np
import cv2
import glob
import os


def visualize(fig, rows, cols, imgs, titles, figname=None, cmap=None):
    for i, img in enumerate(imgs):
        plt.subplot(rows, cols, i+1)
        #plt.title(i+1)
        plt.title(titles[i])
        if (cmap == None):
            img_dims = len(img.shape)
            if img_dims < 3:
                plt.imshow(img, cmap='hot')
            else:
                plt.imshow(img)
        else:
            plt.imshow(img, cmap=cmap)
    plt.show()

    if (figname != None):
        fig.savefig(figname)


# Function that takes an image,
# start and stop positions in both x and y, 
# window size (x and y dimensions),  
# and overlap fraction (for both x and y)
def slide_window(img, x_start_stop=[None, None], y_start_stop=[None, None], 
                    xy_window=(64, 64), xy_overlap=(0.5, 0.5)):
    # If x and/or y start/stop positions not defined, set to image size
    if x_start_stop[0] == None:
        x_start_stop[0] = 0
    if x_start_stop[1] == None:
        x_start_stop[1] = img.shape[1]
    if y_start_stop[0] == None:
        y_start_stop[0] = 0
    if y_start_stop[1] == None:
        y_start_stop[1] = img.shape[0]
    # Compute the span of the region to be searched    
    xspan = x_start_stop[1] - x_start_stop[0]
    yspan = y_start_stop[1] - y_start_stop[0]
    # Compute the number of pixels per step in x/y
    nx_pix_per_step = np.int(xy_window[0]*(1 - xy_overlap[0]))
    ny_pix_per_step = np.int(xy_window[1]*(1 - xy_overlap[1]))
    # Compute the number of windows in x/y
    nx_buffer = np.int(xy_window[0]*(xy_overlap[0]))
    ny_buffer = np.int(xy_window[1]*(xy_overlap[1]))
    nx_windows = np.int((xspan-nx_buffer)/nx_pix_per_step) 
    ny_windows = np.int((yspan-ny_buffer)/ny_pix_per_step) 
    # Initialize a list to append window positions to
    window_list = []
    # Loop through finding x and y window positions
    for ys in range(ny_windows):
        for xs in range(nx_windows):
            # Calculate window position
            startx = xs*nx_pix_per_step + x_start_stop[0]
            endx = startx + xy_window[0]
            starty = ys*ny_pix_per_step + y_start_stop[0]
            endy = starty + xy_window[1]
            
            # Append window position to list
            window_list.append(((startx, starty), (endx, endy)))
    # Return the list of windows
    return window_list

# Define a function to draw bounding boxes
def draw_boxes(img, bboxes, color=(0, 0, 255), thick=6):
    # Make a copy of the image
    imcopy = np.copy(img)
    # Iterate through the bounding boxes
    for bbox in bboxes:
        # Draw a rectangle given bbox coordinates
        cv2.rectangle(imcopy, bbox[0], bbox[1], color, thick)
    # Return the image copy with boxes drawn
    return imcopy
 
# Read car and non-car images from files
def get_training_data(nsamples='ALL'):

    basedir = 'data/vehicles/vehicles/'

    image_types = os.listdir(basedir)
    cars = []

    for imtype in image_types:
        cars.extend(glob.glob(basedir+imtype+'/*'))

    print('Number of Vehicle Images found:',len(cars))

    basedir = 'data/non-vehicles/non-vehicles/'

    image_types = os.listdir(basedir)
    notcars = []

    for imtype in image_types:
        notcars.extend(glob.glob(basedir+imtype+'/*'))

    print('Number of Non-Vehicle Images found:',len(notcars))

    if (nsamples == 'ALL'):
        train_cars = cars
        train_notcars = notcars
    else:
        random_idxs = np.random.randint(0, len(cars), nsamples)

        train_cars = np.array(cars)[random_idxs]
        train_notcars = np.array(notcars)[random_idxs]

    return train_cars, train_notcars
