#!/usr/bin/python

import numpy as np
import cv2
import matplotlib.pyplot as plt
from scipy.misc import imread, imsave

def lab_color_mask(image):
    a_channel_lab = image
    a_thresh = (140, 255)
    a_binary = np.zeros_like(a_channel_lab)
    a_binary[(a_channel_lab >= a_thresh[0]) & (a_channel_lab <= a_thresh[1])] = 1
    a_binary = np.dstack((a_binary, a_binary, a_binary)) * 255
    return a_binary

def to_lab(img_list, label_list):
    lab_list = []
    for i in range(len(img_list)):
        lab_list.append(cv2.cvtColor(img_list[i], cv2.COLOR_RGB2LAB)[:, :, 1])
    return lab_list

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()

def save_image(img, name):
    path = '../output_images/' + name + '.png'
    imsave(path, img)

nimgs = 6
img_list = [imread('../output_images/warped_image0' + 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)]
lab_list = to_lab(img_list, label_list)
lab_bin = lab_color_mask(lab_list)


"""
lab_bin = []
for k in range(len(lab_list)):
    lab_bin.append(lab_color_mask(lab_list[k]))
    save_image(lab_bin[k], label_list[k])

show_images(lab_list, label_list)
#plt.show()
"""