import cv2
import argparse
import enum
import abc
import numpy as np

from board_detction import Board_19x19, VISION_PREPROCESSING_MODE, Board_Detection_SM, detect_board

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

cam = cv2.VideoCapture(args.camera)

WINDOW_NAME = "Output Go Board Detection"
cv2.namedWindow(WINDOW_NAME)
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()

try:
    while True:
        t_start = cv2.getTickCount()
        ret, frame = cam.read()
        width, height = frame.shape[1], frame.shape[0]
        
        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_proc = detect_board(state_machine.state, frame)
        
        board_conture = state_machine.advance_state(board_contures_candiates)
        
        frame_dection = cv2.drawContours(frame.copy(), board_contures_candiates, -1, (0, 255, 0), 1)
        
        frame_classifications = frame.copy()
        
        if board_conture is not None:
            N_line_positions = []
            S_line_positions = []
            W_line_positions = []
            E_line_positions = []
            
            for i in range(19):
                N_line_positions.append(board_messurements.get_cell_position_on_line(i, board_conture[0], board_conture[1]))
                E_line_positions.append(board_messurements.get_cell_position_on_line(i, board_conture[1], board_conture[2]))
                S_line_positions.append(board_messurements.get_cell_position_on_line(i, board_conture[2], board_conture[3]))
                W_line_positions.append(board_messurements.get_cell_position_on_line(i, board_conture[3], board_conture[0]))
            
            N_line_positions = np.array(N_line_positions, dtype=np.int32)
            E_line_positions = np.array(E_line_positions, dtype=np.int32)
            S_line_positions = np.array(list(reversed(S_line_positions)), dtype=np.int32)
            W_line_positions = np.array(list(reversed(W_line_positions)), dtype=np.int32)
        
            cv2.circle(frame_dection, board_conture[0].astype(np.int32), 3, (0, 255, 255), -1)
            cv2.circle(frame_dection, board_conture[1].astype(np.int32), 3, (0, 255, 255), -1)
            cv2.circle(frame_dection, board_conture[2].astype(np.int32), 3, (0, 255, 255), -1)
            cv2.circle(frame_dection, board_conture[3].astype(np.int32), 3, (0, 255, 255), -1)
            
            for p in N_line_positions:
                cv2.circle(frame_dection, p, 2, (0, 0, 255), -1)
            for p in E_line_positions:
                cv2.circle(frame_dection, p, 2, (0, 0, 255), -1)
            for p in S_line_positions:
                cv2.circle(frame_dection, p, 2, (0, 0, 255), -1)
            for p in W_line_positions:
                cv2.circle(frame_dection, p, 2, (0, 0, 255), -1)
                
            for p0, p1 in zip(N_line_positions, S_line_positions):
                cv2.line(frame_dection, tuple(p0), tuple(p1), (0, 0, 255), 1)
                
                for i in range(19):
                    p2 = board_messurements.get_cell_position_on_line(i, p0, p1)
                    cv2.circle(frame_dection, p2, 2, (0, 0, 255), -1)
                
            for p0, p1 in zip(E_line_positions, W_line_positions):
                cv2.line(frame_dection, tuple(p0), tuple(p1), (0, 255, 0), 1)
                
                for i in range(19):
                    p2 = board_messurements.get_cell_position_on_line(i, p0, p1)
                    cv2.circle(frame_dection, p2, 2, (0, 255, 0), -1)
                    
            MESSURE_WIDTH = 15
            STONE_DECTION_WIDTH = 5
            for x in range(19):
                for y in range(19):
                    mid_point = board_messurements.get_cell_position_on_line(x, N_line_positions[y], S_line_positions[y])
                    subsection = frame[mid_point[1] - MESSURE_WIDTH//2:mid_point[1] + MESSURE_WIDTH//2, mid_point[0] - MESSURE_WIDTH//2:mid_point[0] + MESSURE_WIDTH//2]
                    avrage_color = np.mean(subsection, axis=(0, 1)).astype(np.int32).tolist()
                    
                    cv2.rectangle(frame_classifications, (mid_point[0] - MESSURE_WIDTH//2, mid_point[1] - MESSURE_WIDTH//2), (mid_point[0] + MESSURE_WIDTH//2, mid_point[1] + MESSURE_WIDTH//2), avrage_color, -1)
                    cv2.rectangle(frame_classifications, (mid_point[0] - MESSURE_WIDTH//2, mid_point[1] - MESSURE_WIDTH//2), (mid_point[0] + MESSURE_WIDTH//2, mid_point[1] + MESSURE_WIDTH//2), (0, 0, 255), 1)
                    
                    avrage_stone_color = np.mean(frame_edge, axis=(0, 1)).astype(np.int32).tolist()
                    cv2.rectangle(frame_edge_copy, (mid_point[0] - STONE_DECTION_WIDTH//2, mid_point[1] - STONE_DECTION_WIDTH//2), (mid_point[0] + STONE_DECTION_WIDTH//2, mid_point[1] + STONE_DECTION_WIDTH//2), (avrage_stone_color, avrage_stone_color, avrage_stone_color), -1)
                    cv2.rectangle(frame_edge_copy, (mid_point[0] - STONE_DECTION_WIDTH//2, mid_point[1] - STONE_DECTION_WIDTH//2), (mid_point[0] + STONE_DECTION_WIDTH//2, mid_point[1] + STONE_DECTION_WIDTH//2), (0, 0, 255), 1)
                    
        res = np.concatenate((
            np.concatenate((
                frame,
                frame_classifications,
                frame_edge_copy
            ), axis=1),
            np.concatenate((
                cv2.cvtColor(frame_proc, cv2.COLOR_GRAY2BGR),
                frame_dection,
                frame
            ), 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)
        
        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...")
            break
    
    
except KeyboardInterrupt:
    print("Exiting")

cam.release()
cv2.destroyAllWindows()