import cv2
import numpy as np
from ultralytics import YOLO as YOLOv8
import PySimpleGUI as sg
from pathlib import Path

def RemoveBlackBorder(image, shrink=True):
    image = np.array(image)
    
    copyImg = cv2.cvtColor(image.copy(), cv2.COLOR_BGR2HSV_FULL)
    mask = np.zeros(image[:,:,0].shape, dtype=np.uint8)
    mask[(copyImg[:,:,1] > 35)] = 255
    copyImg = cv2.cvtColor(copyImg, cv2.COLOR_HSV2RGB_FULL)
    resROI = cv2.bitwise_and(copyImg, copyImg, mask=mask)
    
    image_gray = cv2.cvtColor(resROI, cv2.COLOR_RGB2GRAY)
    _, thresh = cv2.threshold(image_gray, 0, 255, cv2.THRESH_BINARY)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15))
    morph = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
    contours = cv2.findContours(morph, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    contours = contours[0] if len(contours) == 2 else contours[1]
    bigCont = max(contours, key=cv2.contourArea)
    x, y, w, h = cv2.boundingRect(bigCont)
    sfx, sfy =  (w *0.01, h*0.07) if shrink else (0, 0)
    crop = image[y +int(sfy) : y + h - int(sfy), x + int(sfx): x + w - int(sfx)]
    return crop

def RemoveBlackBorder2(image):
    image = np.array(image)
    
    copyImg = cv2.cvtColor(image.copy(), cv2.COLOR_BGR2HSV)
    h = copyImg[:,:,0]
    mask = np.ones(h.shape, dtype=np.uint8) * 255
    th = (25, 175)
    mask[(h > th[0]) & (h < th[1])] = 0
    copyImg = cv2.cvtColor(copyImg, cv2.COLOR_HSV2BGR)
    resROI = cv2.bitwise_and(copyImg, copyImg, mask=mask)
        
    image_gray = cv2.cvtColor(resROI, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(image_gray, 0, 255, cv2.THRESH_BINARY)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15))
    morph = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
    contours = cv2.findContours(morph, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    contours = contours[0] if len(contours) == 2 else contours[1]
    bigCont = max(contours, key=cv2.contourArea)
    x, y, w, h = cv2.boundingRect(bigCont)
    crop = image[y : y + h, x: x + w ]
    return crop

def LoadAndVeiwImge (Path:str, YOLO:float = None):
    image = np.load(Path)
    image = image.astype(np.uint8)
    
    print(image.shape)
    
    if YOLO > 0:
        YoloModel = YOLOv8(model="RARP_YoloV8_ROI.pt")
        res = YoloModel.predict(image, stream=True)
        
        for r in res:
            conf = r.boxes.conf.cpu().numpy()
            for i, box in enumerate(r.boxes.xyxy.cpu().numpy()):
                if conf[i] > YOLO:
                    box = box.astype(int)
                    print (i, box, conf[i])   
                    #x, y = (box[0], box[1])
                    #xw, yh = (box[2], box[3])
                    cv2.rectangle(image, (box[0],box[1]), (box[2],box[3]), (0, 255, 150*i), 3)
    
    cv2.namedWindow("Output", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("Output", (1024, 720))
    cv2.imshow("Output", image)
    cv2.waitKey(0)
    cv2.destroyAllWindows() 

if __name__ == "__main__":
    sg.theme('LightBlue')
    
    layout = [
        [sg.Text('Image Viewer', size=(40, 1), justification='center', font='Helvetica 20',key='-status-')],
        [sg.Text('Camera number: ', size=(15, 1)), sg.InputText(default_text='',  size=(80, 1),key='-Path-')],
        [sg.Text('YOLO Conf.: ', size=(15, 1)), sg.InputText(default_text='0.40',  size=(20, 1),key='-YOLO-')],
        [sg.Checkbox("RAW Image", default=True, key="-NP_Array_IMG-"), sg.Checkbox("Remove Black Bars Old ver.", default=False, key="-oldVer-")],
        [sg.Image(filename='', key='image')],
        [
            sg.Button('YOLO', size=(10, 1), font='Helvetica 14',key ='-start-'),
            sg.Button('Remove BB', size=(10, 1), font='Helvetica 14',key = '-Remove_BB-'),
            sg.Button('View', size=(10, 1), font='Helvetica 14',key = '-verImg-'),
            sg.Button('Exit', size=(10, 1), font='Helvetica 14', key='-exit-'), 
        ]
    ]
    
    window = sg.Window('Image Viewer',layout, location=(100, 100))
    recording = False
    
    while True:
        try:
            event, vals = window.read(timeout=20)
            match event:
                case None | "-exit-":
                    break
                case "-start-":
                    if vals["-NP_Array_IMG-"]:
                        vals["-YOLO-"] = 0 if vals["-YOLO-"] == "" else vals["-YOLO-"]
                        LoadAndVeiwImge(vals["-Path-"], float(vals["-YOLO-"]))
                    else:
                        sg.popup_error("Only RAW images are accepted *.npy")
                case "-Remove_BB-":
                    image = np.load(vals["-Path-"]) if vals["-NP_Array_IMG-"] else cv2.imread(vals["-Path-"], cv2.IMREAD_COLOR)
                    if vals["-oldVer-"] :
                        image = RemoveBlackBorder(image, False)
                    else:
                        image = RemoveBlackBorder2 (image)
                    cv2.namedWindow("Output", cv2.WINDOW_NORMAL)
                    cv2.resizeWindow("Output", (1024, 720))
                    cv2.imshow("Output", image)
                    cv2.waitKey(0)
                    cv2.destroyAllWindows() 
                case "-verImg-":
                    image = np.load(Path(vals["-Path-"])) if vals["-NP_Array_IMG-"] else cv2.imread(vals["-Path-"], cv2.IMREAD_COLOR)
                    cv2.namedWindow("Output", cv2.WINDOW_NORMAL)
                    cv2.resizeWindow("Output", (1024, 720))
                    cv2.imshow("Output", image)
                    cv2.waitKey(0)
                    cv2.destroyAllWindows() 
                    
        except Exception as e:
            print(e)
            
                
    window.close()