import cv2
import numpy as np

"""
Hodnoty treba upravit, podla robota. Najlepsie je ich zistit podla edit_contrast.py
"""
CONTRAST_VALUE = 2
SATURATION_VALUE = 2.950
THRESHOLD_VALUE = 235

MARGIN_H = 0
MARGIN_S = 0
MARGIN_U = 0

Letter_H_Contour = np.load("ReferenceData/contours/Image_H_Contour.npy")
Letter_S_Contour = np.load("ReferenceData/contours/Image_S_Contour.npy")
Letter_U_Contour = np.load("ReferenceData/contours/Image_U_Contour.npy")

def get_victim_on_wall(camera):
    """
    Urci pismeno alebo farbu na stene pomocou kamery.
    ### Parametre:
        - kamera z ktorej chceme urcit pismeno
    ### Vracia:
        - Pismeno ktore tam je z [H,S,U] alebo farbu z [Red, Green, Yellow]
        - Ak tam ziadne nie je vrati "None", ak error tak "Error"
    """

    is_frame, colour_frame = camera.read()

    if not is_frame:
        return "Could not get frame from camera"

    hsv = cv2.cvtColor(colour_frame, cv2.COLOR_BGR2HSV)
    hsv[:, :, 1] = np.clip(hsv[:, :, 1] * SATURATION_VALUE, 0, 255)
    bgr_scaled = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)

    contrasted = cv2.convertScaleAbs(bgr_scaled, alpha=float(CONTRAST_VALUE), beta=0)
    gray_scale_frame = cv2.cvtColor(contrasted, cv2.COLOR_BGR2GRAY)

    _, frame_thresh = cv2.threshold(gray_scale_frame, int(THRESHOLD_VALUE), 255, cv2.THRESH_BINARY_INV)

    frame_contours, _ = cv2.findContours(frame_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    if not frame_contours:
        return "None"
    
    biggest_contour = max(frame_contours, key=cv2.contourArea)

    mask = np.zeros(gray_scale_frame.shape, dtype=np.uint8)

    cv2.drawContours(mask, [biggest_contour], -1, 255, thickness=cv2.FILLED)

    mean_b, mean_g, mean_r, _ = cv2.mean(colour_frame, mask=mask)

    average_b = int(mean_b)
    average_g = int(mean_g)
    average_r = int(mean_r)

    #Hodnoty je potrebne upravit podla robota
    if (average_r < 70) and (average_g > 80) and (average_b < 90):
        return "Green"
    elif (average_r > 80) and (average_g < 80) and (average_b < 80):
        return "Red"    
    elif (average_r > 100) and (average_g > 100) and (average_b < 80):
        return "Yellow"

    s_score = cv2.matchShapes(biggest_contour, Letter_S_Contour, cv2.CONTOURS_MATCH_I2, 0.0)
    h_score = cv2.matchShapes(biggest_contour, Letter_H_Contour, cv2.CONTOURS_MATCH_I2, 0.0)
    u_score = cv2.matchShapes(biggest_contour, Letter_U_Contour, cv2.CONTOURS_MATCH_I2, 0.0)

    if s_score < MARGIN_S:
        return "S"
    elif h_score < MARGIN_H:
        return "H", h_score
    elif u_score < MARGIN_U:
        return "U"
    else:
        return "None"
