import functools
import cv2
import numpy as np

# ============================================================================

def determine_threshold(image):
    window_main = "Determine Threshold"
    cv2.namedWindow(window_main)

    result = { 'threshold' : 127 , 'mask' : None }
    def on_trackbar_threshold(threshold_value):
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        _, mask = cv2.threshold(gray, threshold_value, 255, cv2.THRESH_BINARY_INV)

        display_image = np.hstack([image
            , cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
            , cv2.bitwise_and(image, image, mask=mask)
        ])
        cv2.imshow(window_main, display_image)

        result['threshold'] = threshold_value
        result['mask'] = mask

    trackbar_thresh = "Threshold"
    on_trackbar_threshold(result['threshold'])
    cv2.createTrackbar(trackbar_thresh, window_main, result['threshold'], 255, on_trackbar_threshold)
    on_trackbar_threshold(cv2.getTrackbarPos(trackbar_thresh, window_main))

    cv2.waitKey(0)
    cv2.destroyWindow(window_main)
    return result

# ============================================================================

def extract_clean_tile(image, mask, background_color):
    window_main = "Extract Clean Tile"
    cv2.namedWindow(window_main)

    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    sorted_contours = sorted(contours, key=lambda x: cv2.contourArea(x), reverse=True)
    contour = sorted_contours[0]

    masked_image = cv2.bitwise_and(image, image, mask=mask) # Cache for display

    result = { 'tile' : None, 'mask' : None }
    def on_trackbar_outline(scale):
        approx_contour = cv2.approxPolyDP(contour, (scale / 20.0), True)
        x, y, w, h = cv2.boundingRect(approx_contour)
        PAD_X, PAD_Y = 8, 8
        tx, ty, tw, th = x - PAD_X, y - PAD_Y , w + PAD_X * 2, h + PAD_Y * 2

        output_image_a = masked_image.copy()
        cv2.drawContours(output_image_a, [contour], -1, (0,0,255), 1)

        cleaned_image = image.copy()
        cv2.drawContours(cleaned_image, [approx_contour], -1, background_color, 3, cv2.LINE_AA)

        cleaned_mask = np.zeros_like(mask)
        cv2.drawContours(cleaned_mask, [approx_contour], -1, (255,255,255), -1)
        cleaned_masked_image = cv2.bitwise_and(cleaned_image, cleaned_image, mask=cleaned_mask)

        output_image_b = cleaned_masked_image.copy()
        cv2.rectangle(output_image_b, (x - 1, y - 1), (x + w, y + h), (0, 191,191), 1)

        display_image = np.hstack([masked_image, output_image_a, output_image_b])
        cv2.imshow(window_main, display_image)

        result['tile'] = cleaned_masked_image[y:y+h,x:x+w]
        result['mask'] = cleaned_mask[y:y+h,x:x+w]

    trackbar_outline = "Approx."
    on_trackbar_outline(0)
    cv2.createTrackbar(trackbar_outline, window_main, 50, 200, on_trackbar_outline)
    on_trackbar_outline(cv2.getTrackbarPos(trackbar_outline, window_main))

    cv2.waitKey(0)
    cv2.destroyWindow(window_main)
    return result

# ============================================================================

# NB: expects BGR tile
def tile_pattern(tile, mask, rows, cols, background_color):
    window_main = "Tile Pattern"
    cv2.namedWindow(window_main)

    trackbar_csx, trackbar_csy, trackbar_rsx, trackbar_rsy = "col_step_x", "col_step_y", "row_step_x", "row_step_y"

    tile_height, tile_width = tile.shape[:2]
    context = { trackbar_csx : tile_width
        , trackbar_csy : 0
        , trackbar_rsx : 0
        , trackbar_rsy : tile_height
    }

    def update_display():
        col_step_x, col_step_y = context[trackbar_csx], context[trackbar_csy]
        row_step_x, row_step_y = context[trackbar_rsx], context[trackbar_rsy]
        #canvas_height = tile_height + row_step_y * (rows - 1) + col_step_y * (cols - 1)
        #canvas_width = tile_width + col_step_x * (cols - 1) + row_step_x * (rows - 1)
        canvas_height, canvas_width = tile_height * (rows + 1), tile_width * (cols + 1)
        canvas = np.zeros((canvas_height, canvas_width, 3),  np.uint8)
        canvas = canvas + np.uint8(background_color)

        output_mask = np.repeat((mask != 0)[:, :, np.newaxis], 3, axis=2)

        start_pos_x, start_pos_y = 0, 0
        for x in range(rows):
            pos_x, pos_y = start_pos_x, start_pos_y
            for y in range(cols):
                dest = canvas[pos_y:pos_y+tile_height, pos_x:pos_x+tile_width]
                np.putmask(dest, output_mask, tile)
                pos_x, pos_y = pos_x + row_step_x, pos_y + row_step_y
                if pos_x > start_pos_x + tile_width:
                    pos_x -= tile_width
            start_pos_x, start_pos_y = start_pos_x + col_step_x, start_pos_y + col_step_y
            if start_pos_y > tile_height:
                start_pos_y -= tile_height
        cv2.imshow(window_main, canvas)

    def on_trackbar_step(scale, target):
        context[target] = scale
        update_display()

    update_display()

    cv2.createTrackbar(trackbar_csx, window_main, context[trackbar_csx], tile_width
        , functools.partial(on_trackbar_step, target=trackbar_csx))
    cv2.createTrackbar(trackbar_csy, window_main, context[trackbar_csy], tile_height
        , functools.partial(on_trackbar_step, target=trackbar_csy))
    cv2.createTrackbar(trackbar_rsx, window_main, context[trackbar_rsx], tile_width
        , functools.partial(on_trackbar_step, target=trackbar_rsx))
    cv2.createTrackbar(trackbar_rsy, window_main, context[trackbar_rsy], tile_height
        , functools.partial(on_trackbar_step, target=trackbar_rsy))

    cv2.waitKey(0)
    cv2.destroyWindow(window_main)

# ============================================================================

tile_image = cv2.imread('tile05.png')
BACKGROUND_COLOR = (0, 0, 0)

thresh_result = determine_threshold(tile_image)

tile_result = extract_clean_tile(tile_image, thresh_result['mask'], BACKGROUND_COLOR)

tile_pattern(tile_result['tile'], tile_result['mask'], 3, 3, BACKGROUND_COLOR)