import cv2
import argparse
import enum
import uuid
import numpy as np

from model import load_model, TARGET_SIZE, classify_stone_empty


# Parse arguments
parser = argparse.ArgumentParser()
parser.add_argument("camera", type=int, help="Camera index")
args = parser.parse_args()

cam = cv2.VideoCapture(args.camera)

stone_detection_model = load_model("models/model1.pth")

WINDOW_NAME = "Output Go Board Detection"
cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL)
cv2.resizeWindow(WINDOW_NAME, 1920, 1080)
cv2.moveWindow(WINDOW_NAME, 1,0)

font = cv2.FONT_HERSHEY_SIMPLEX
font_color = (127, 255, 0)

fps = 0

board_messurements = Board_19x19()
state_machine = Board_Detection_SM()


fourcc = cv2.VideoWriter_fourcc(*"mp4v")
recording = False

try:
    while True:
        t_start = cv2.getTickCount()
        ret, frame = cam.read()
        

        frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        frame_edge = cv2.GaussianBlur(frame_gray, (5, 5), 0)
        frame_edge = cv2.adaptiveThreshold(frame_edge, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 9, 2)
        frame_edge_copy = cv2.cvtColor(frame_edge.copy(), cv2.COLOR_GRAY2BGR)

        board_contures_candiates, frame_edge_processed = detect_board(state_machine.state, frame)

        #board_conture = state_machine.advance_state(board_contures_candiates)


        if len(board_contures_candiates) > 0:
            board_conture = np.squeeze(board_contures_candiates[0], axis=1)
            frame_dection = cv2.drawContours(frame.copy(), board_contures_candiates, -1, (0, 255, 0), 1)
            frame_classifications = frame.copy()
        else:
            board_conture = None
            frame_dection = frame.copy()
            frame_classifications = frame.copy()

        # Stage 2: Classfiy Stones
        if board_conture is not None:
            
            # Getting Perspectiv correct Image
            board_image = corret_board_perspective(frame, board_conture, s=400)
            display_board_image = board_image.copy()
            draw_board_markers(display_board_image, np.array([
                [0, 0],
                [0, board_image.shape[1]],
                [board_image.shape[0], board_image.shape[1]],
                [board_image.shape[0], 0]    
            ]), board_messurements)
            
            display_board_image = pad_center(display_board_image, frame.shape[0:2]) 
            board_masked, board_antimasked1, board_antimasked2, board_pre_processed = preprocess_for_stone_dection(board_image)
            
            classification_output = board_image.copy()
            classification_color_output = board_image.copy()
            
            pre_classifcation = classfy_board(board_masked, board_antimasked1, board_antimasked2, classification_output, board_messurements, windows_size=10)
            color_classifcation = classfiy_color(board_image, pre_classifcation, classification_color_output, board_messurements, windows_size=10)
            
            display_classifcation_preprocessed = pad_center(cv2.cvtColor(board_pre_processed, cv2.COLOR_GRAY2BGR), frame.shape[0:2])
            display_classifcation_output = pad_center(classification_output, frame.shape[0:2])
            display_classifcation_color_output = pad_center(classification_color_output, frame.shape[0:2])

        else:
            display_board_image = np.full(frame.shape, 0, dtype=np.uint8)
            display_classifcation_preprocessed = np.full(frame.shape, 0, dtype=np.uint8)
            display_classifcation_output = np.full(frame.shape, 0, dtype=np.uint8)
            display_classifcation_color_output = np.full(frame.shape, 0, dtype=np.uint8)
            display_classifcation_preprocessed = np.full(frame.shape, 0, dtype=np.uint8)

        if recording:
            VideoOut.write(frame)

        # Added Labeling
        frame_edge_processed = cv2.cvtColor(frame_edge_processed, cv2.COLOR_GRAY2BGR)
        cv2.putText(frame, "Raw Input Frame", (10, frame.shape[0]-10), font, 1, font_color, 2, cv2.LINE_AA)   
        cv2.putText(frame_edge_processed, "Board Edges", (10, frame.shape[0]-10), font, 1, font_color, 2, cv2.LINE_AA)   
        cv2.putText(display_board_image, "Board (Perspecive Corrected)", (10, frame.shape[0]-10), font, 1, font_color, 2, cv2.LINE_AA)   
        cv2.putText(display_classifcation_output, "Stone Pre-Classfication", (10, frame.shape[0]-10), font, 1, font_color, 2, cv2.LINE_AA)   
        cv2.putText(display_classifcation_color_output, "Stone Classfication", (10, frame.shape[0]-10), font, 1, font_color, 2, cv2.LINE_AA)   
        cv2.putText(display_classifcation_preprocessed, "Classfication Preprocessed", (10, frame.shape[0]-10), font, 1, font_color, 2, cv2.LINE_AA)   
    
        res = np.concatenate((
            np.concatenate((
                frame,
                display_board_image,
                display_classifcation_preprocessed,
            ), axis=1),
            np.concatenate((
                frame_edge_processed,
                display_classifcation_output,
                display_classifcation_color_output,
            ), axis=1),
            np.concatenate((
                frame_edge_processed,
                display_classifcation_output,
                display_classifcation_color_output,
            ), axis=1)
        ), axis=0)
        
        for n,t in enumerate([
                f"FPS: {fps:.2f}",
                f"GoBoard Dections: {len(board_contures_candiates)}", 
            ] + state_machine.create_info_text()):
            res = cv2.putText(res, t, (10, 30 + 15*n), font, 0.5, font_color, 1, cv2.LINE_AA)
        
        if recording:
            cv2.circle(res, (20, 20), 10, (0, 0, 255), -1)
        
        cv2.imshow(WINDOW_NAME, res)
        
        t_end = cv2.getTickCount()
        
        fps = cv2.getTickFrequency() / (t_end - t_start)
        
        k = cv2.waitKey(1)
        if k%256 == 27: # ESC pressed
            print("Escape hit, closing...")
            break0
        elif k%256 == 32:
            cv2.imwrite("frame.png", frame)
            
        elif k%256 == 114:# # R-Key
            if recording:
                recording = False
                VideoOut.release()
            else:
                recording = True
                VideoOut = cv2.VideoWriter(f"vid_{str(uuid.uuid4())}.mp4", fourcc, 30.0, (frame.shape[1], frame.shape[0]))
    
    
except KeyboardInterrupt:
    print("Exiting")

cam.release()
cv2.destroyAllWindows()