import sys
import base64

import cv2
import ast
from collections import Counter
from sklearn.cluster import KMeans
import numpy as np
import pandas as pd
import pickle as cPickle
import face_cropper


def extract_face_images(image):
    face_model = face_cropper.FaceCropper(min_face_detector_confidence=0.5, face_detector_model_selection=0,
                                          landmark_detector_static_image_mode=True,
                                          min_landmark_detector_confidence=0.5)
    face_image, hair_image, eyes = face_model.get_head_segments(image)
    return face_image, hair_image, eyes


def palette_perc(k_cluster):
    n_pixels = len(k_cluster.labels_)
    counter = Counter(k_cluster.labels_)  # count how many pixels per cluster
    perc = {}
    for i in counter:
        perc[i] = np.round(counter[i] / n_pixels, 2)
    perc = dict(sorted(perc.items()))

    palette = pd.DataFrame()
    for i in range(0, 5):
        add = pd.DataFrame({"color": str(list(k_cluster.cluster_centers_[i])), "perc": perc[i]}, index=[0])
        palette = pd.concat([palette, add], axis=0, ignore_index=True)
    palette.color = palette.color.apply(lambda x: ast.literal_eval(x))

    for i in range(0, len(palette)):
        color_list = palette.color[i]
        j2 = list(filter(lambda x: x >= 230, color_list))
        if len(j2) == 3:
            palette = palette.drop(i, axis=0)
    palette = palette.reset_index(drop=True)
    return palette


def resize_image(img):
    wanted_width = 500
    height, width, _ = img.shape
    aspect_ratio = width / height
    wanted_height = int(wanted_width / aspect_ratio)
    resized_img = cv2.resize(img, (wanted_width, wanted_height))
    return resized_img


def find_color_palette(img):
    img = resize_image(img)
    clt = KMeans(n_clusters=5, n_init=10)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    clt_1 = clt.fit(img.reshape(-1, 3))
    palette = palette_perc(clt_1)
    selected_color = (palette[palette.perc == palette.perc.max()].reset_index(drop=True)).loc[0, "color"]
    return selected_color


def color_season(image):
    try:
        face_image, hair_image, eyes = extract_face_images(image)
    except:
        return ""
    hair_color = find_color_palette(hair_image)
    skin_color = find_color_palette(face_image)
    eye_color = find_color_palette(eyes)

    with open('./color_palette_classifier.pkl', 'rb') as file_id:
        classifier = cPickle.load(file_id)

    personal_colors = pd.DataFrame({"hair_r": hair_color[0], "hair_g": hair_color[1], "hair_b": hair_color[2],
                                    "eye_r": eye_color[0], "eye_g": eye_color[1], "eye_b": eye_color[2],
                                    "skin_r": skin_color[0], "skin_g": skin_color[1], "skin_b": skin_color[2]},
                                   index=[0])
    selected_palette = classifier.predict(personal_colors)
    return selected_palette[0].replace("_", " ")


def color_analysis(im_b64):
    im_bytes = base64.b64decode(im_b64)
    im_arr = np.frombuffer(im_bytes, dtype=np.uint8)  # im_arr is one-dim Numpy array
    image = cv2.imdecode(im_arr, flags=cv2.IMREAD_COLOR)
    color_theme = color_season(image)

    if color_theme == "soft spring":
        color_theme = "light spring"
    if color_theme == "clear summer":
        color_theme = "soft summer"
    if color_theme == "clear autumn":
        color_theme = "deep autumn"
    if color_theme == "clear winter":
        color_theme = "deep winter"
    if color_theme == "soft winter":
        color_theme = "clear winter"
    return color_theme


if __name__ == "__main__":
    im_b64 = sys.argv[1]
    print(color_analysis(im_b64))