import cv2
import numpy as np

GREEN = cv2.imread("ReferenceData/Images/Image_GreekOmega.jpg")
RED = cv2.imread("ReferenceData/Images/Image_GreekPhi.jpg")
YELLOW = cv2.imread("ReferenceData/Images/Image_GreekPsi.jpg")
IH = cv2.imread("ReferenceData/Images/Image_H.jpg")
IS = cv2.imread("ReferenceData/Images/Image_S.jpg")
IU = cv2.imread("ReferenceData/Images/Image_U.jpg")


def proc(img, a, s):
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    hsv[:, :, 1] = np.clip(hsv[:, :, 1] * float(s), 0, 255)
    bgr = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
    # return a color BGR image (apply contrast but keep color)
    return cv2.convertScaleAbs(bgr, alpha=float(a), beta=0)

cv2.namedWindow("Contrast Sliders", cv2.WINDOW_NORMAL)
cv2.createTrackbar("Contrast", "Contrast Sliders", 100, 500, lambda v: None)
cv2.createTrackbar("Saturation", "Contrast Sliders", 100, 500, lambda v: None)
cv2.createTrackbar("Threshold", "Contrast Sliders", 100, 255, lambda v: None)

final_alpha = None
while True:
    a = cv2.getTrackbarPos("Contrast", "Contrast Sliders") / 100.0
    s = cv2.getTrackbarPos("Saturation", "Contrast Sliders") / 100.0
    imgs1 = [proc(GREEN, a, s), proc(RED, a, s), proc(YELLOW, a, s)]
    imgs2 = [proc(IH, a, s), proc(IS, a, s), proc(IU, a, s)]

    # thresholded rows for both sets
    t = cv2.getTrackbarPos("Threshold", "Contrast Sliders")
    thr_imgs1 = []
    thr_imgs2 = []
    for im in imgs1:
        gray = im if im.ndim == 2 else cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
        _, thr = cv2.threshold(gray, t, 255, cv2.THRESH_BINARY_INV)
        thr_color = cv2.cvtColor(thr, cv2.COLOR_GRAY2BGR)
        thr_imgs1.append(thr_color)
    for im in imgs2:
        gray = im if im.ndim == 2 else cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
        _, thr = cv2.threshold(gray, t, 255, cv2.THRESH_BINARY_INV)
        thr_color = cv2.cvtColor(thr, cv2.COLOR_GRAY2BGR)
        thr_imgs2.append(thr_color)

    top_row1 = cv2.hconcat(imgs1)
    bottom_row1 = cv2.hconcat(thr_imgs1)
    col1 = cv2.vconcat([top_row1, bottom_row1])

    top_row2 = cv2.hconcat(imgs2)
    bottom_row2 = cv2.hconcat(thr_imgs2)
    col2 = cv2.vconcat([top_row2, bottom_row2])

    # place columns side-by-side
    canvas = cv2.hconcat([col1, col2])
    cv2.imshow("Contrast Sliders", canvas)
    k = cv2.waitKey(50) & 0xFF
    if k == ord('e'):
        v = float(input("Alpha: "))
        pos = int(round(v * 100.0))
        if 100 < pos < 500:
            cv2.setTrackbarPos("Contrast", "Contrast Sliders", pos)
        continue
    if k in (27, ord('q')):
        final_alpha = a
        final_saturation = s
        final_threshold = t
        break
cv2.destroyAllWindows()
print(f"contrast={final_alpha:.3f}, saturation={final_saturation:.3f}, threshold={final_threshold}")