import numpy as np
import cv2
import json
import matplotlib.pyplot as plt
import matplotlib.pylab as pyl

import cStringIO
from PIL import Image
import urllib

def load_image(url):
    '''Read remote images, so we don't need to download the files
    Input: URL string
    Otput: RGB image as a numpy array (width, height, 3)'''
    f = Image.open(cStringIO.StringIO(urllib.urlopen(url).read()))
    return np.array(f)

def gray_image(image):
    '''Transform to gray scale'''
    return cv2.cvtColor(image, cv2.cv.CV_BGR2GRAY) #openCV read the color values as BGR instead of RGB that uses matplotlib

def trained_classifier(file):
    '''Get the already trained classifier from an XML file
    Input:  path to the XML file
    Output: haar-cascade classifier'''
    return cv2.CascadeClassifier(file)
    

def face_detector(image, classifier, classifier1):
    '''Run the two classifiers against the provided image
    Input:  image, trained classifier
    Output: numpy n-dimensional array with the detected faces
    every face is defined by the array (x,y,width,height)'''
    
    faces1 = classifier.detectMultiScale(image, scaleFactor=1.1, minNeighbors=3)
    faces2 = classifier1.detectMultiScale(image, scaleFactor=1.1, minNeighbors=3)
    if len(faces1) > 0 and len(faces2) > 0:
        return np.concatenate((faces1, faces2), axis=0)
    elif len(faces1) == 0:
        return faces2
    else:
        return faces1

    
def area(face, img_area):
    '''Calculate the area of a face. Used as sorted key function on select_face method.
    If the face area is bigger than 1/6 of the photo, discard it (selfie or wrong detection)
    Input: numpy array with a face
    Output: (int) face area'''
    result = 0
    if type(face) != int:
        face_area = face[2] * face[3]
        if face_area < (1/6) * img_area:
            result = face_area
    return result


def select_face(faces, img_area):
    '''Select the biggest face on the image
    Input: numpy n-dim array with the faces
    Output: numpy array with the selected face'''
    if len(faces) > 0:
        face_areas = {}
        for face in faces:
            face_areas[area(face, img_area)] = face
        return face_areas[max(face_areas.keys())]   
    else:
        return 0
    
    
def draw_rects(img, rect):
    '''Draw a rectangle on the image'''
    x,y,w,h = rect
    return cv2.rectangle(img, (x, y), (x+w, y+h), (0,255,0), 2)

#Use two different trained sets for the classifier
cascade_classifier =  trained_classifier('/opencv-2.4.8/data/haarcascades/haarcascade_frontalface_default.xml')
cascade_classifier1 =  trained_classifier('/opencv-2.4.8/data/haarcascades/haarcascade_frontalface_alt.xml')

def face_detection(image_url):
    '''Input: (string) Image URL
    Output: (tuple) (numpy array RGB image, numpy array gray image, numpy array with face coordinates)'''
    img = load_image(image_url)
    gray = gray_image(img)
    faces = face_detector(gray, cascade_classifier, cascade_classifier1)
    selected_face = select_face(faces, img.shape[0] * img.shape[1])
    return (img, gray, selected_face)

def crop_image(image, gray, face):
    '''Crop the image by selecting the ROI
    Input: image, gray_image, array of faces
    Output: (ROI cropped RGB, gray cropped)'''
    im = image.copy()
    im1 = gray.copy()
    yt, xt, rgb = image.shape
    x,y,w,h = face
    #wide: select 1.5 * w/2 to each side
    #height: 12*heads or until the end of the image
    tall = 12*h if 10*h < yt else yt
    middle = x + w/2
    start_w = int(round(middle - 1.5*w)) if middle - 1.5*w > 0 else 0
    end_w = start_w + 3*w
    return (im[y+h:tall,start_w:end_w], im1[y+h:tall,start_w:end_w])

#Load some images as examples:
filename = '../text/boards/wendy_s_lookbook.json'
with open(filename, "r") as f:
    pins = json.loads(f.read())
pins = pins[:30]

images = [face_detection(pin['pin_image']) for pin in pins]
images = [(img, gray, face) for img, gray, face in images if face is not 0]
copies = [(image.copy(), face) for image, gray, face in images]
rects = [draw_rects(img, face) for img, face in copies] 
cropped = [crop_image(img,gray,face) for img,gray,face in images]

for i, image in enumerate(images):
    print 'Face detection and ROI selection of image ' + str(i)
    fig = plt.figure(i, figsize=(20,4))
    plt.subplot(141),plt.imshow(image[0])
    plt.title('Input Image'), plt.xticks([]), plt.yticks([])
    plt.subplot(142)
    plt.hist(image[1].ravel(),255,[0,256])
    plt.title('Histogram of gray image')
    plt.subplot(143), plt.imshow(copies[i][0])
    plt.title('Selected face'), plt.xticks([]), plt.yticks([])
    plt.subplot(144), plt.imshow(cropped[i][0])
    plt.title('ROI Image'), plt.xticks([]), plt.yticks([])
    plt.show()

def create_mask(gray_cropped):
    '''Create a new image as mask
    Input: gray image with the ROI area
    Output: numpy array with the mask'''
    
    #canny edge detection, threshold values: 20, 100
    edges = cv2.Canny(gray_cropped,20,100)
    #gaussian filter
    blur = cv2.GaussianBlur(edges,(5,5),0)
    #closing transformation with a kernel of 2x2 size
    kernel = np.ones((2,2),np.uint8)
    closing = cv2.morphologyEx(blur,cv2.MORPH_CLOSE,kernel, iterations = 2)
    #dilation with kernel 2x2
    dilation = cv2.dilate(closing,kernel,iterations=5)

    # find contours
    dil = dilation.copy() #make a copy as it gets modified by findContours
    contours,hierarchy = cv2.findContours(dil,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
    # find largest contour
    areas = [cv2.contourArea(cnt) for cnt in contours]
    max_cont = contours[np.argmax(areas)]
    
    # draw largest contour on the mask
    mask = np.zeros((gray_cropped.shape[0], gray_cropped.shape[1]),np.uint8)
    cv2.drawContours(mask, [max_cont], 0, 255, -1)
    return mask

masks = [create_mask(gray) for img,gray in cropped]

#results = [images[i][0][masks[i] == 255] for i in range(len(masks))]
#For example purposes, paint the background in fluor green
results = []
for i in range(len(masks)):
    cp = cropped[i][0].copy()
    cp[masks[i] == 0] = [204, 255, 0]
    results.append(cp)

for i, image in enumerate(cropped):
    print 'Background removal ' + str(i)
    fig = plt.figure(i, figsize=(20,4))
    plt.subplot(131),plt.imshow(cropped[i][0])
    plt.title('Cropped Image'), plt.xticks([]), plt.yticks([])
    plt.subplot(132), plt.imshow(masks[i], 'gray')
    plt.title('Mask'), plt.xticks([]), plt.yticks([])
    plt.subplot(133), plt.imshow(results[i])
    plt.title('Result on image'), plt.xticks([]), plt.yticks([])
    plt.show()

import scipy as sp
from sklearn.cluster import DBSCAN
from scipy.cluster.vq import kmeans,vq

def perform_dbscan(image, eps):
    '''Apply DBSCAN algorithm
    Input: image, eps (max distance between pixels of same cluster)
    Output: clustered row image, labeled row image'''
    rows, cols, rgb = image.shape
    image_row = np.reshape(image,(rows * cols, 3))
    
    #number of min_samples: 2% of total pixels of image
    num_samples = round(rows * cols * 0.02)
    
    db = DBSCAN(eps=eps, min_samples=num_samples).fit(image_row)
    core_samples = db.core_sample_indices_
    labels = db.labels_
    # Number of clusters in labels, ignoring noise if present (label -1)
    n_clusters_ = len(set(labels)) - (1 if -1 in labels else 0)
    
    clusters = [image_row[labels == i] for i in xrange(n_clusters_)]
    return clusters, labels

def perform_kmeans_dbscan(image):
    '''Apply kmeans algorithm to select the central color
    Input:  row image
    Output: centroid RGB color'''
    centroid,_ = kmeans(image, 1)
    return centroid

def color_detection_method1(image, eps=15):
    '''Apply dbscan and kmeans
    Input: image
    Output: clustered image'''
    clusters, labels = perform_dbscan(image, eps)
    clus = []
    for cluster in clusters: #obtain the main color per each cluster detected by dbscan
        clus.append(perform_kmeans_dbscan(cluster))
    image1 = image.copy()
    rows, cols, rgb = image1.shape
    image_row = np.reshape(image1,(rows * cols, 3))
    for i in xrange(len(clus)):
        image_row[labels == i] = clus[i][0] #assign to each pixel the value of the centroid color
    image_row[labels == -1 ] = [255, 0, 255] #display the outliers with magenta color
    return np.reshape(image_row,(rows, cols, 3))

def perform_kmeans(image, k):
    '''Apply kmeans algorithm.
    Input:  image, number of clusters
    Output: centroids, vector with the index of the centroid assigned to each pixel, clusters'''
    rows, cols, rgb = image.shape
    # reshape the image in a single row array of RGB pixels
    image_row = np.reshape(image,(rows * cols, 3))
    # perform the clustering
    centroids,_ = kmeans(image_row, k)
    #vector quantization, assign to each pixel the index of the nearest centroid (k=1..4)
    qnt,_ = vq(image_row,centroids)
    # reshape the qnt vector to the original image shape
    image_centers_id = np.reshape(qnt,(rows, cols))
    # assign the centroid value to each pixel
    clustered = centroids[image_centers_id]
    return centroids, qnt, clustered

dbscan = [color_detection_method1(image) for image in results]

kmeans = [perform_kmeans(image, 5) for image in results]

#plot the results
for i, image in enumerate(results):
    print 'Clustering colors on image: ' + str(i)
    
    fig = plt.figure(i, figsize=(15,5))
    plt.subplot(131), plt.imshow(image)
    plt.title("Cropped with mask"), plt.xticks([]), plt.yticks([])
    plt.subplot(132), plt.imshow(dbscan[i])
    plt.title("Clustered image with DBSCAN + Kmeans"), plt.xticks([]), plt.yticks([])
    plt.subplot(133), plt.imshow(kmeans[i][2])
    plt.title("Clustered image with Kmeans (k=4)"), plt.xticks([]), plt.yticks([])
    plt.show()

import webcolors
from colormath.color_objects import LabColor, sRGBColor
from colormath.color_conversions import convert_color

def colors_to_hex(color):
    '''Transform a RGB color to its hex value
    Input: array of RGB values
    Output: array of hex values'''
    return webcolors.rgb_to_hex(color)

def hex_to_rgb(hex_color):
    '''Transform Hex color to RGB'''
    return webcolors.hex_to_rgb(hex_color)

def hex_to_lab(hex_color):
    '''Transform Hex color to Lab'''
    rgb = webcolors.hex_to_rgb(hex_color)
    return convert_color(sRGBColor(rgb[0], rgb[1], rgb[2]), LabColor)

#load the lab_matrix and lab_colors pickle files as numpy arrays
import pickle

lab_matrix_file = open('lab-matrix.pk', 'rb')
lab_colors_file = open('lab-colors.pk', 'rb')

lab_matrix = pickle.load(lab_matrix_file)
lab_colors = pickle.load(lab_colors_file)

from vectorize_delta import delta_e_cie2000 #vectorized version

def closest_color(color, lab_matrix, lab_colors):
    '''Select the Lab color with less deltaE CIE 2000 distance to color
    from the 197525 different Lab colors on the lab_matrix np.array
    Input: hex color, lab_matrix, lab_colors
    Output: (string) color name out of 17 palette on lab_colors'''
    color = hex_to_lab(color)
    dist = delta_e_cie2000(np.array([color.lab_l, color.lab_a, color.lab_b]), lab_matrix)
    min_col = lab_matrix[np.argmin(dist)]
    return lab_colors[np.argmin(dist)]

palette_colors = {'black': '#000000', 'blue': '#0000ff', 'navy': '#000080', 'white': '#ffffff', 'beige': '#f5f5dc', 
                  'red': '#ff0000', 'brown': '#a52a2a', 'gold': '#ffd700', 'silver': '#c0c0c0', 'gray': '#808080', 
                  'pink': '#ffc0cb', 'green': '#008000', 'yellow': '#ffff00', 'purple': '#800080', 'orange': '#ffa500'}

def rgb_to_palette(hex_color):
    '''Assign to an RGB color the closest CSS color by selecting the one with minimum Euclidian distance
    Input: Hex color
    Output: Palette color name'''
    rgb_color = hex_to_rgb(hex_color)
    min_colors = {}
    for name, hex_val in palette_colors.items():
        r_prim, g_prim, b_prim = webcolors.name_to_rgb(name)
        rd = (r_prim - rgb_color[0]) ** 2
        gd = (g_prim - rgb_color[1]) ** 2
        bd = (b_prim - rgb_color[2]) ** 2
        min_colors[(rd + gd + bd)] = name
    return min_colors[min(min_colors.keys())]

def count_colors(qnt, centroids):
    '''Count the number of pixels associated to each cluster
    Input: quantization vector, centroids
    Output: k length array with the count of pixels per cluster''' 
    counts, bins = sp.histogram(qnt, len(centroids))
    return counts, bins

def count_cluster_colors(centroids, qnt, clustered):
    '''Count the number of pixels per centroid (only needed for displaying the results)
    Ugly trick to not consider the green background color in the results'''
    colors = []
    centroids1 = []
    for i, centroid in enumerate(centroids):
        if centroid[0] != 204 and centroid[1] != 255 and centroid[2] != 0:
            centroids1.append(centroid)
            colors.append(colors_to_hex(centroid))
    counts, bins = count_colors(qnt, centroids1)
    return counts, colors, clustered

def print_results(img, mask, colors, pixels_per_cluster, clustered, num):
    '''Plot the results: 
    Histogram of the clustered colors by kmeans, detected face, mask and clustered image'''
    fig = plt.figure(num, figsize=(18,5))
    ax = plt.subplot(141)
    width = 0.8
    barlist = ax.bar(range(len(pixels_per_cluster)), pixels_per_cluster)
    for i in xrange(len(colors)):
        barlist[i].set_color(colors[i])
    ax.set_xticks(np.arange(len(pixels_per_cluster)) + width/2)
    ax.set_xticklabels(colors, rotation=90)
    ax.set_title("       #Pixels associated with each cluster")
    plt.ylabel("#Pixels")
    plt.xlabel("Centroid colors (hex value)")
    plt.subplot(142), plt.imshow(img)
    plt.title("Selected face"), plt.xticks([]), plt.yticks([])
    plt.subplot(143), plt.imshow(mask)
    plt.title("Mask"), plt.xticks([]), plt.yticks([])
    plt.subplot(144), plt.imshow(clustered)
    plt.title("Clustered"), plt.xticks([]), plt.yticks([])
    plt.show()
    fig.savefig(str(num) + '.png', dpi = (200))

cluster_counts = [count_cluster_colors(centroids, qnt, clustered) for centroids, qnt, clustered in kmeans]

for i in range(len(results)):
    print_results(copies[i][0], results[i], cluster_counts[i][1], cluster_counts[i][0], cluster_counts[i][2], i)
    print 'Color discretization:'
    print 'deltaE CIE2000 approach: ' +  str([closest_color(color, lab_matrix, lab_colors) for color in cluster_counts[i][1]])
    print 'RGB euclidian distance: ' + str([rgb_to_palette(rgb_color) for rgb_color in cluster_counts[i][1]])
    print