import cv2 as cv
import numpy as np

from util.image_processing import raster_image

blue_lower_hsv = np.array([90, 75, 75])
blue_upper_hsv = np.array([130, 255, 255])

red_lower1_hsv = np.array([0, 75, 75])
red_upper1_hsv = np.array([10, 255, 255])
red_lower2_hsv = np.array([170, 75, 75])
red_upper2_hsv = np.array([179, 255, 255])

yellow_lower_hsv = np.array([15, 75, 75])
yellow_upper_hsv = np.array([40, 255, 255])

white_lower_hsv = np.array([0, 0, 190])
white_upper_hsv = np.array([180, 65, 255])

black_lower_hsv = np.array([0, 0, 1]) # because 0 0 0 pixels are ignored as background
black_upper_hsv = np.array([180, 255, 30])

def get_color_percentage(image: cv.Mat) -> dict:
	hsv = cv.cvtColor(image, cv.COLOR_BGR2HSV)

	mask_red1 = cv.inRange(hsv, red_lower1_hsv, red_upper1_hsv)
	mask_red2 = cv.inRange(hsv, red_lower2_hsv, red_upper2_hsv)
	mask_red = mask_red1 | mask_red2
	mask_blue = cv.inRange(hsv, blue_lower_hsv, blue_upper_hsv)
	mask_yellow = cv.inRange(hsv, yellow_lower_hsv, yellow_upper_hsv)
	mask_white = cv.inRange(hsv, white_lower_hsv, white_upper_hsv)
	mask_black = cv.inRange(hsv, black_lower_hsv, black_upper_hsv)

	width, height = image.shape[:2]
	image_pixels = width*height

	red_pixels = cv.countNonZero(mask_red) 
	blue_pixels = cv.countNonZero(mask_blue) 
	yellow_pixels = cv.countNonZero(mask_yellow) 
	white_pixels = cv.countNonZero(mask_white)
	black_pixels = cv.countNonZero(mask_black)
	
	red_percentage = int((red_pixels / image_pixels) * 100)
	blue_percentage = int((blue_pixels / image_pixels) * 100)
	yellow_percentage = int((yellow_pixels / image_pixels) * 100)
	white_percentage = int((white_pixels / image_pixels) * 100)
	black_percentage = int((black_pixels / image_pixels) * 100)
	none_percentage = 100 - red_percentage - blue_percentage - yellow_percentage - white_percentage - black_percentage

	return {
		"red": red_percentage,
		"blue": blue_percentage,
		"yellow": yellow_percentage,
		"white": white_percentage,
		"black": black_percentage,
		"none": none_percentage
	}

def get_raster_color_percentage(image: cv.Mat, rows: int, cols: int) -> list[list[dict]]:
	rasters2d = raster_image(image, cols, rows)
	ret = []
	for raster1d in rasters2d:
		new = []
		for raster in raster1d:
			color_percentage = get_color_percentage(raster)
			new.append(color_percentage)
		ret.append(new)
	return ret

def get_edges(gray_image: cv.Mat):
    edges = cv.Canny(cv.blur(gray_image, (5,5)), threshold1=100, threshold2=200)
    edge_count = np.count_nonzero(edges)
    return edge_count, edges

def get_corners(gray_image: cv.Mat):
    max_corners  = 100
    quality_level = 0.14
    min_distance = 1
    corners = cv.goodFeaturesToTrack(cv.blur(gray_image, (5,5)), maxCorners=max_corners, qualityLevel=quality_level, minDistance=min_distance)
    if corners is not None:
        corners = np.intp(corners) # int0 not working for me :(
        return len(corners)
    else:
        return 0
    
def get_contours(edged_gray_image: cv.Mat):
    contours, _ = cv.findContours(edged_gray_image, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
    contour_count = len(contours)
    return contour_count
