import cv2
import argparse
import enum
import abc
from typing import Optional
import numpy as np

class Board_Sizes_ABC(abc.ABC):
    
    @abc.abstractmethod
    def get_cell_rel_pos(self, n: int):
        pass
    
    @abc.abstractmethod
    def get_num_cells(self) -> tuple[int, int]:
        pass
    
    def get_cell_position_on_line(self, n: int, p0: np.ndarray, p1: np.ndarray):
        vec = p1 - p0
        return np.round(p0 + vec*self.get_cell_rel_pos(n)).astype(np.int32)

    def iterat_pos(self) -> tuple[tuple[float, float], tuple[int, int]]:
        for x in range(self.N + 1):
            for y in range(self.N + 1):
                yield (x, y), (self.get_cell_rel_pos(x), self.get_cell_rel_pos(y))

class Board_19x19(Board_Sizes_ABC):
    LINE_WIDTH = 1
    CELL_SIZE = 21
    PADDING = 11.5
    N = 18
    
    TOTAL_WIDTH = N*CELL_SIZE + (N+1)*LINE_WIDTH + 2*PADDING
    
    def __init__(self):
        self.rel_lut = [(self.PADDING + self.LINE_WIDTH*0.5 + (self.LINE_WIDTH + self.CELL_SIZE)*i)/self.TOTAL_WIDTH for i in range(19)]
    
    def get_num_cells(self):
        return (self.N+1, self.N+1)
    
    def get_cell_rel_pos(self, n: int):
        return self.rel_lut[n]

class VISION_PREPROCESSING_MODE(enum.Enum):
    CANNY = "CANNY"
    THRES_HOLD = "THRES_HOLD"
    THRES = "THRES"
    TEST_CANNY_HOLD = "TEST_CANNY_HOLD"
    TEST_CANNY = "TEST_CANNY"

class Board_Detection_SM:
    
    def __init__(self):
        self.state = VISION_PREPROCESSING_MODE.CANNY
        self.hold_time = 0
        self.dection_history = []
        self.jitter = 0
        self.num_decetions = 0
        self.num_double_decetions = 0
        self.last_decetion_time = 0
        self.last_double_detection_time = 0

    def advance_state(self, detections: list[np.array]) -> Optional[np.array]:
        detection = self.process_nun_decetions(detections)
        
        match self.state:
            case VISION_PREPROCESSING_MODE.CANNY:
                if self.last_decetion_time > 50:
                    self.state = VISION_PREPROCESSING_MODE.THRES_HOLD
                    self.hold_time = 100
            
            case VISION_PREPROCESSING_MODE.THRES_HOLD:
                if self.hold_time == 0:
                    self.state = VISION_PREPROCESSING_MODE.THRES
                else:
                    self.hold_time -= 1
            
            case VISION_PREPROCESSING_MODE.THRES:
                if self.last_decetion_time > 50 or self.jitter > 3 or self.num_double_decetions > 15:
                    self.state = VISION_PREPROCESSING_MODE.TEST_CANNY_HOLD
                    self.hold_time = 200
            
            case VISION_PREPROCESSING_MODE.TEST_CANNY_HOLD:
                if self.hold_time == 0:
                    self.state = VISION_PREPROCESSING_MODE.TEST_CANNY
                else:
                    self.hold_time -= 1
                    
            case VISION_PREPROCESSING_MODE.TEST_CANNY:
                if self.last_decetion_time < 10 and self.jitter < 3:
                    self.state = VISION_PREPROCESSING_MODE.CANNY
                else:
                    self.state = VISION_PREPROCESSING_MODE.THRES_HOLD
                    self.hold_time = 100
            
            case _:
                raise ValueError("Invalid preprocessing mode")
            
        return detection
      
    def process_nun_decetions(self, detections: list[np.array]) -> Optional[np.array]:
        num_decetions = len(detections)
        dection = None
        
        if num_decetions > 2:
            self.num_double_decetions += 1
            self.num_last_double_detection_time = 0
        else:
            if self.last_double_detection_time > 300:
                self.num_double_decetions = 0
            else:
                self.last_double_detection_time = max(1000, self.last_double_detection_time)
                
        if num_decetions == 1:
            self.last_decetion_time = 0
            
            self.dection_history.append(np.squeeze(detections[0], axis=1))
            self.dection_history = self.dection_history[-10:]
        else:
            self.last_decetion_time += 1
        
        if len(self.dection_history) > 0:
            dection_history_np = np.array(self.dection_history)
            avrage_detection = np.mean(dection_history_np, axis=0)
            dection = avrage_detection.copy()
        
        if len(self.dection_history) > 6:
            self.jitter = 0
            for i in dection_history_np:
                self.jitter = np.max(np.abs(i - avrage_detection))
                
        return dection     
    
    def create_info_text(self) -> list[str]:
        
        return [
            f"Preprocessing: {self.state.value}",
            f"Double Dection: {self.num_double_decetions}",
            f"Jitter: {self.jitter:.2f}",
            f"Last Dection Time: {self.last_decetion_time}",
            f"Last Double Dection Time: {self.last_double_detection_time}",
            f"Hold Time: {self.hold_time}"
        ]
        
    
    
    
def detect_board(state: VISION_PREPROCESSING_MODE, frame: np.ndarray) -> tuple[list[np.array], np.ndarray]:
    
    frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        
    if state == VISION_PREPROCESSING_MODE.CANNY or \
        state == VISION_PREPROCESSING_MODE.TEST_CANNY or \
        state == VISION_PREPROCESSING_MODE.TEST_CANNY_HOLD:
            
        frame_proc = cv2.GaussianBlur(frame_gray, (3, 3), 0)
        frame_proc = cv2.Canny(frame_proc, 130, 200)
        
    elif state == VISION_PREPROCESSING_MODE.THRES_HOLD or state == VISION_PREPROCESSING_MODE.THRES:
        frame_proc = cv2.GaussianBlur(frame_gray, (3, 3), 0)
        frame_proc = cv2.adaptiveThreshold(frame_proc, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 19, 3)
    
    else:
        raise ValueError("Invalid preprocessing mode")
    
    countours, _ = cv2.findContours(frame_proc, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    return list(
        filter(lambda x: type(x) == np.ndarray,
            map(filter_countours, countours)
        )   
    ), frame_proc
        
def filter_countours(c: np.array) -> Optional[np.array]:
    
    esp = cv2.arcLength(c, True)*0.05
    approx = cv2.approxPolyDP(c, esp, True)
    
    # Stage 1: Has 4 corners
    if len(approx) != 4:
        return None
    
    # Stage 2: Area is big enough
    area = cv2.contourArea(approx) 
    
    if area < 1000:
        return None
    
    # Stage 3: Are all side findContourslengths similar
    lengths = np.array([
        np.linalg.norm(approx[0] - approx[1]),
        np.linalg.norm(approx[1] - approx[2]),
        np.linalg.norm(approx[2] - approx[3]),
        np.linalg.norm(approx[3] - approx[0])
    ])
    
    avrg_length = np.mean(lengths)
    lengths_diff = np.abs(lengths)/avrg_length
                
    if not (np.all(lengths_diff < 1.1) and np.all(lengths_diff > 0.90)):
        return None
    
    return approx

def draw_board_markers(f: np.array, conture: np.array, board_messurements: Board_Sizes_ABC): 
    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, conture[0], conture[1]))
        E_line_positions.append(board_messurements.get_cell_position_on_line(i, conture[1], conture[2]))
        S_line_positions.append(board_messurements.get_cell_position_on_line(i, conture[2], conture[3]))
        W_line_positions.append(board_messurements.get_cell_position_on_line(i, conture[3], 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(f, conture[0].astype(np.int32), 3, (0, 255, 255), -1)
    cv2.circle(f, conture[1].astype(np.int32), 3, (0, 255, 255), -1)
    cv2.circle(f, conture[2].astype(np.int32), 3, (0, 255, 255), -1)
    cv2.circle(f, conture[3].astype(np.int32), 3, (0, 255, 255), -1)
    
    for p in N_line_positions:
        cv2.circle(f, p, 2, (0, 0, 255), -1)
    for p in E_line_positions:
        cv2.circle(f, p, 2, (0, 0, 255), -1)
    for p in S_line_positions:
        cv2.circle(f, p, 2, (0, 0, 255), -1)
    for p in W_line_positions:
        cv2.circle(f, p, 2, (0, 0, 255), -1)
        
    for p0, p1 in zip(N_line_positions, S_line_positions):
        cv2.line(f, 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(f, p2, 2, (0, 0, 255), -1)
        
    for p0, p1 in zip(E_line_positions, W_line_positions):
        cv2.line(f, 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(f, p2, 2, (0, 255, 0), -1)
            
    cv2.drawContours(f, [conture], -1, (0, 255, 255), 1)
    
def corret_board_perspective(frame: np.array, conture: np.array, s=250):
    p1 = conture.astype(np.float32)
    p2 = np.array([
        [s,0], # Top right
        [0, 0], # Top Left
        [0,s], # Bottom Right
        [s,s], # Bottom Left
    ], dtype=np.float32)
    
    T = cv2.getPerspectiveTransform(p1, p2)
    res = cv2.warpPerspective(frame, T, (s, s))
    res = res[0:s,0:s]

    return res

def pad_center(frame: np.array, target_size: np.array):
    f = frame[0:target_size[1],0:target_size[0]]
    
    horiztontal_padding = (target_size[1]-f.shape[1]) // 2
    vertical_padding = (target_size[0]-f.shape[0]) // 2
    
    return cv2.copyMakeBorder(
        f, 
        vertical_padding, vertical_padding, 
        horiztontal_padding, horiztontal_padding, 
        cv2.BORDER_CONSTANT, value=(0, 0, 0)
    )