import sys

import cv2
import numpy as np


def main(fn):
    SZ = 21
    OUTPUT_ARRAY = [[0]*SZ for i in range(SZ)]
    img = cv2.imread(fn)
    gr = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gr = cv2.threshold(gr, 190, 255, cv2.THRESH_BINARY_INV)[1]

    kernel = np.ones((2, 2), dtype=np.uint8)
    gr = cv2.morphologyEx(gr, cv2.MORPH_OPEN, kernel)
    gr = cv2.morphologyEx(gr, cv2.MORPH_CLOSE, kernel)

    (__im2, cnts, __hierarchy) = cv2.findContours(gr, cv2.RETR_EXTERNAL,
                                                  cv2.CHAIN_APPROX_SIMPLE)

    # get the biggest shape
    cnts.sort(key=cv2.contourArea)
    cnt = cnts[-1]
    cv2.drawContours(img, [cnt], -1, (0, 255, 255), 4)

    # find the enclosing square
    x, y, w, h = cv2.boundingRect(cnt)
    cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)

    mask = np.zeros(gr.shape, dtype=np.uint8)

    # now, we know that the field is 31x31
    w0 = w / SZ
    h0 = h / SZ
    for i in range(SZ):
        for j in range(SZ):
            cv2.rectangle(
                mask,
                (int(x + i * w0 + 10), int(y + j * w0 + 12)),
                (int(x + i * w0 + w0 - 14), int(y + j * h0 + h0 - 12)),
                255, -1,
            )

    gr = cv2.bitwise_and(gr, mask)

    KNOWN_NUMBERS = [
        # number, x-coord, y-coord
        (1, 9, 2), (1, 10, 16), (1, 14, 12),

        (2, 9, 5), (2, 2, 6), (2, 2, 7), (2, 12, 17), (2, 9, 15),

        (3, 10, 1), (3, 10, 11), (3, 10, 14), (3, 10, 18), (3, 18, 5),
        (3, 2, 4), (3, 16, 9),

        (4, 10, 2), (4, 1, 7), (4, 1, 8),

        (5, 11, 2), (5, 18, 9), (5, 1, 11), (5, 3, 16),
        (6, 13, 1), (6, 9, 16),

        (7, 13, 2), (7, 1, 10), (7, 3, 4), (7, 11, 0),

        (8, 9, 1), (8, 6, 4), (8, 20, 11), (8, 4, 6),

        (9, 6, 3), (9, 6, 12),
    ]

    #serious_mark(gr, KNOWN_NUMBERS, 1, x, y, w0, h0)

    TEMPLATES = []
    for (n, xx, yy) in KNOWN_NUMBERS:
        TEMPLATES.append((n, get_sq_function(gr, xx, yy, x, y, w0, h0)))

    # visualize_search(gr, TEMPLATES, x, y, w0, h0, 10, 16)

    out = cv2.cvtColor(gr, cv2.COLOR_GRAY2RGB)
    for i in range(SZ):
        for j in range(SZ):
            sq = get_sq_function(gr, i, j, x, y, w0, h0)
            if sq is None:
                continue

            results = []
            for (number, templ) in TEMPLATES:
                result = cv2.matchTemplate(sq, templ, cv2.TM_CCOEFF)
                (_, score, _, _) = cv2.minMaxLoc(result)
                if score > 1:
                    results.append((score, number))
            if results:
                results.sort()
                mark_square(out, i, j, x, y, w0, h0, results[-1][1])
                OUTPUT_ARRAY[j][i] = results[-1][1]

    cv2.imwrite("out.jpg", out)
    return OUTPUT_ARRAY


def serious_mark(img, nums, hl, x, y, w0, h0):
    for (n, xx, yy) in nums:
        if n != hl:
            continue
        mark_square(img, xx, yy, x, y, w0, h0, n)
    cv2.imwrite("out.jpg", img)
    sys.exit(0)


def visualize_search(img, TEMPLATES, x, y, w0, h0, xx, yy):
    out = np.zeros((300, 2000), dtype=np.uint8)
    sq = get_sq_function(img, xx, yy, x, y, w0, h0)

    X = 20
    results = []
    for (number, templ) in TEMPLATES:
        result = cv2.matchTemplate(sq, templ, cv2.TM_CCOEFF)
        out[20:20 + sq.shape[0], X:X + sq.shape[1]] = sq
        out[150:150 + templ.shape[0], X:X + templ.shape[1]] = templ
        (_, score, _, _) = cv2.minMaxLoc(result)
        results.append((score, number))
        X += sq.shape[1] + 5

    best = max([i[0] for i in results])
    X = 20
    for (val, num) in results:
        cv2.putText(out, "%0.2f" % (val / best, ),
                    (X, 120),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.4, 255, 1, cv2.LINE_AA)
        X += sq.shape[1] + 5

    cv2.imwrite("out.jpg", out)
    sys.exit(0)


def find_fit_square(img, x, y, x0, y0, w0, h0):
    sub_img = extract_square(img, x, y, x0, y0, w0, h0)
    (__im2, cnts, __hierarchy) = cv2.findContours(sub_img, cv2.RETR_EXTERNAL,
                                                  cv2.CHAIN_APPROX_SIMPLE)

    if not cnts:
        return None

    min_x = min_y = 100000
    max_x = max_y = -100000
    for i in cnts:
        c = i.reshape((i.shape[0] * i.shape[-1]))
        min_x = min((min_x, min(c[0::2])))
        max_x = max((max_x, max(c[0::2])))
        min_y = min((min_y, min(c[1::2])))
        max_y = max((max_y, max(c[1::2])))

    sub_sub = sub_img[min_y:max_y, min_x:max_y]
    res = cv2.resize(sub_sub, (int(w0), int(h0)), interpolation=cv2.INTER_CUBIC)
    return res


def extract_square(img, x, y, x0, y0, w0, h0):
    x_base = int(x0 + x * w0)
    y_base = int(y0 + y * h0)
    h0 = int(h0)
    w0 = int(w0)
    return img[y_base + 12:y_base + h0 - 12, x_base + 10:x_base + w0 - 14]


get_sq_function = find_fit_square


def mark_square(img, x, y, x0, y0, w0, h0, text='X'):
    COLOR = 255
    if 3 == len(img.shape):
        COLOR = (0, 0, 255)
    cv2.putText(img, str(text),
                (int(x0 + x * w0 + 0.3 * w0), int(y0 + y * h0 + 0.7 * h0)),
                cv2.FONT_HERSHEY_SIMPLEX,
                1.4, COLOR, 2, cv2.LINE_AA)
    cv2.putText(img, "%s" % (x, ),
                (int(x0 + x * w0 + 3), int(y0 + y * h0 + 0.6 * h0)),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.4, COLOR, 1, cv2.LINE_AA)
    cv2.putText(img, "%s" % (y, ),
                (int(x0 + x * w0 + 3), int(y0 + y * h0 + 0.8 * h0)),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.4, COLOR, 1, cv2.LINE_AA)


if "__main__" == __name__:
    ar = main("hlavolam.jpg")
    print(ar)