import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colors, cm
from mpl_toolkits.mplot3d import Axes3D
from scipy.interpolate import splprep, splev

pixel_thresold_pct = 0.05

def get_img_histogram(img, colors, gray):
    """Saves color channel masked images to numbered jpg images."""
    img_n = 0
    for color in colors[0]:
        split_img = img.copy()
        split_img[np.where(gray != color)] = 0
        cv.imwrite(str(img_n) + ".jpg", split_img)
        img_n += 1
    plt.hist(gray.ravel(), 256, [0,256])
    plt.savefig('histogram.jpg')
    plt.show()

def smooth_contours(contours, img):
    smoooth = []
    for contour in contours:
        x, y = contour.T
        x = x.tolist()[0]
        y = y.tolist()[0]
        tck, u = splprep([x, y], u=None, k=1, s=1.0, per=1)
        u_new = np.linspace(u.min(), u.max(), 25)
        x_new, y_new = splev(u_new, tck, der=0)
        res_array = [[[int(i[0]), int(i[1])]] for i in zip(x_new, y_new)]
        smoooth.append(np.asarray(res_array, dtype=np.int32))
    cv2.drawContours(img, smoooth, -1, (255,255,255), 2)
    copy = img.copy()
    cv.imwrite('contour_smooth.jpg', copy)

def color_range_mask(hsv, color_start, color_end, name):
    mask = cv.inRange(hsv, color_start, color_end)
    result = cv.bitwise_and(img, img, mask=mask)
    result[mask == 255] = (255, 255, 255)
    cv.imwrite(f'{name}.png', result)
    return result

def get_pixel_representations(img):
    # Get color representations from rgb
    rgb = cv.cvtColor(img, cv.COLOR_BGR2RGB)
    pixel_colors = rgb.reshape((np.shape(rgb)[0] * np.shape(rgb)[1], 3))
    norm = colors.Normalize(vmin=-1., vmax=1.)
    norm.autoscale(pixel_colors)
    pixel_colors = norm(pixel_colors).tolist()
    return pixel_colors

def threshold_and_contour(image, name):
    gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
    ret, thresh = cv.threshold(gray, 150, 255, cv.THRESH_BINARY)
    cv.imwrite(f'thresholded_{name}.png', thresh)
    contours, _ = cv.findContours(image=thresh, mode=cv.RETR_TREE, method=cv.CHAIN_APPROX_NONE)
    # Smooth contours
    copy = img.copy()
    cv.drawContours(image=copy, contours=contours, contourIdx=-1, color=(0, 255, 0), thickness=2, lineType=cv.LINE_AA)
    cv.imwrite(f'contoured_{name}.png', copy)


def get_hsv_plot(img, pixel_colors):
    # Get HSV color representation
    hsv = cv.cvtColor(img, cv.COLOR_BGR2HSV)
    h, s, v = cv.split(hsv)
    h_flat = h.flatten()
    s_flat = s.flatten()
    v_flat = v.flatten()
    fig = plt.figure()
    axis = fig.add_subplot(1, 1, 1, projection="3d")
    max_value = np.max(v_flat)
    # Filter for coords TODO: make this better
    coords = []
    for idx, value in enumerate(v_flat):
        if value >= 0.99 * max_value:
            coords.append((h_flat[idx], s_flat[idx], value))
    coords = list(set(coords))
    axis.scatter(h_flat, s_flat, v_flat, facecolors=pixel_colors, marker=".")
    axis.set_xlabel("Hue")
    axis.set_ylabel("Saturation")
    axis.set_zlabel("Value")
    plt.savefig('hsv_spectrum.jpg')
    plt.close()
    return hsv

if __name__ == '__main__':
    # Blur and resize original image
    img = cv.imread("cowboy_bebop_test.jpg")
    blur = cv.blur(img,(4, 4))
    blur0=cv.medianBlur(blur,5)
    blur1 = cv.GaussianBlur(blur, (5,5),0)
    cv.imwrite('blur.jpg', blur1)

    pixel_colors = get_pixel_representations(img)
    hsv = get_hsv_plot(blur1, pixel_colors)

    # Two manually selected colors based on hsv plot
    color_start_o = (0, 60, 100)
    color_end_o = (70, 305, 255)

    color_start_b = (0, 0, 0)
    color_end_b = (300, 300, 50)

    mask1 = color_range_mask(hsv, color_start_o, color_end_o, 'mask1')
    mask2 = color_range_mask(hsv, color_start_b, color_end_b, 'mask2')

    threshold_and_contour(mask1, 'color_mask1')
    threshold_and_contour(mask2, 'color_mask2')