import cv2
import numpy as np
import operator
import torch



def find_corners(polygon, limit_func, compare_func):
    """
    Input: Rectangle puzzle extract from contours
    Output: One of four cornet point depend on limit_func, compare_func
    # limit_fn is the min or max function
    # compare_fn is the np.add or np.subtract function
    Note: (0,0) point is at the top-left

    top-left: (x+y) min
    top-right: (x-y) max
    bot-left: (x-y) min
    bot-right: (x+y) max
    """

    index, _ = limit_func(enumerate([compare_func(ptr[0][0], ptr[0][1]) for ptr in polygon]), key = operator.itemgetter(1))

    return polygon[index][0][0], polygon[index][0][1]



def draw_circle_at_corners(original, ptr):
    """
    Helper function to draw circle at corners
    """

    cv2.circle(original, ptr, 5, (0,255,0), cv2.FILLED)



def grid_line_helper(img, shape_location, length = 10):
    """
    Helper function to fine vertical, horizontal line
    Find horizontal line: shape_location = 0
    Find vertical line: shape_location = 1
    """

    clone_img = img.copy()
    row_or_col = clone_img.shape[shape_location]

    # Find the distance between lines
    size = row_or_col // length

    # Morphological transfromation to find line

    # Define morphology kernel
    if shape_location == 0:
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (size,1))
    else:
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1,size))

    clone_img = cv2.erode(clone_img, kernel)
    clone_img = cv2.dilate(clone_img, kernel)

    return clone_img



def draw_line(img, lines):
    """
    Draw all lines in lines got from cv2.HoughLine()
    """
    clone_img = img.copy()
    # lines list from cv2.HoughLine() is 3d array
    # Convert to 2d array

    lines = np.squeeze(lines)
    for rho, theta in lines:
        a = np.cos(theta)
        b = np.sin(theta)
        x0 = a*rho
        y0 = b*rho
        x1 = int(x0 + 1000 * (-b))
        y1 = int(y0 + 1000 * a)
        x2 = int(x0 - 1000 * (-b))
        y2 = int(y0 - 1000 * a)
        #Draw line every loop
        cv2.line(clone_img, (x1,y1), (x2,y2), (255,255,255), 4)
    return clone_img



def clean_square_helper(img):
    """
    Clean noises in every square splited
    Input: One of 81 squares
    Output: Cleaned square and boolean var which so that there is number in it
    """

    if np.isclose(img, 0).sum() / (img.shape[0] * img.shape[1]) >= 0.96:
        return np.zeros_like(img), False

    # if there is very little white in the region around the center, this means we got an edge accidently
    height, width = img.shape
    mid = width // 2
    if np.isclose(img[:, int(mid - width * 0.38):int(mid + width * 0.38)], 0).sum() / (2 * width * 0.38 * height) >= 0.98:
        return np.zeros_like(img), False

    # center image
    contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    contours = sorted(contours, key=cv2.contourArea, reverse=True)
    x, y, w, h = cv2.boundingRect(contours[0])

    start_x = (width - w) // 2
    start_y = (height - h) // 2
    new_img = np.zeros_like(img)
    new_img[start_y:start_y + h, start_x:start_x + w] = img[y:y + h, x:x + w]

    return new_img, True



def resize_square(clean_square_list):
    """
    Resize clean squares into 28x28 in order to feed to classifier
    """

    resized_list = []
    for img in clean_square_list:
        resized = cv2.resize(img, (28,28), interpolation=cv2.INTER_AREA)
        resized_list.append(resized)
    return resized_list



def resize_square32(clean_square_list):
    """
    Resize clean squares into 32x32 in order to feed to tf classifier
    """
    resized_list = []
    for img in clean_square_list:
        resized = cv2.resize(img, (32,32), interpolation=cv2.INTER_AREA)
        resized_list.append(resized)
    return resized_list


def classify_one_digit(model, resize_square, org_image):
    """
    Determine whether each square has number by counting number of (not black) pixel and compare to threshold value
    Using classifier to predict number in square
    - Return 0 if the square is blank
    - Return predict digit if the square has number
    """
    threshold = 0
    if (org_image.shape[0] > 600 or org_image.shape[1] > 600) or (org_image.shape[1] > 600 or org_image.shape[2] > 600):
        threshold = 40
    else: 
        threshold = 60
    # Determine blank square

    if (resize_square != resize_square.min()).sum() < threshold:
        return str(0)
    
    model.eval()
    # Convert to shape (1,1,28,28) to be compatible with dataloader for evaluation
    iin = torch.Tensor(resize_square).unsqueeze(0).unsqueeze(0)

    with torch.no_grad():
        out = model(iin)
        # Get index of predict digit
        _, index = torch.max(out, 1)

    pred_digit = index.item()

    return str(pred_digit)



def normalize(resized_list):
    """
    Scale pixel value for recognition
    """

    return [img/255 for img in resized_list]



def convert_str_to_board(string, step = 9):
    """
    Convert recognized string into 2D array for sudoku solving
    """
    
    board = []
    for i in range(0, len(string), step):
        board.append([int(char) for char in string[i:i+step]])
    return np.array(board)