import cv2
import numpy
import pytesseract
from PIL import Image, ImageDraw, ImageFilter

INPUT_IMAGE = 'Unpainted.png'
OUTPUT_IMAGE = 'Painted.png'
SQUARES_W = 31
SQUARES_H = 50

colors = [
    (253, 83, 83, 255),
    (252, 216, 0, 255),
    (170, 170, 170, 255),
    (94, 2, 5, 255),
    (255, 185, 63, 255),
    (251, 2, 0, 255),
    (255, 223, 172, 255),
    (0, 175, 82, 255),
]

pytesseract.pytesseract.tesseract_cmd = 'C:\\Program Files\\Tesseract-OCR\\tesseract.exe'

img = Image.open(INPUT_IMAGE)
w, h = img.size

img = img.crop((67, 104, w - 64, h - 102))
draw = ImageDraw.Draw(img)
square_w = img.size[0] * 1.0 / SQUARES_W
square_h = img.size[1] * 1.0 / SQUARES_H

prev_digit = 2

for y in range(SQUARES_H):
    for x in range(SQUARES_W):
        print("{}%".format((y * SQUARES_W + x) * 100.0 / (SQUARES_H * SQUARES_W)))

        x_start = int(round(x * square_w))
        x_end = int(round((x + 1) * square_w))
        y_start = int(round(y * square_h))
        y_end = int(round((y + 1) * square_h))

        digit_img = img.crop((x_start, y_start, x_end, y_end))
        digit_img = digit_img.filter(ImageFilter.SMOOTH)
        digit_img = digit_img.filter(ImageFilter.SHARPEN)
        digit_img = digit_img.filter(ImageFilter.SMOOTH)
        digit_img = digit_img.filter(ImageFilter.EDGE_ENHANCE)
        digit_img = digit_img.point(lambda p: 255 if p > 200 else p * 0.7)
        # digit_img = digit_img.filter(ImageFilter.SMOOTH)
        # digit_img = digit_img.filter(ImageFilter.EDGE_ENHANCE)
        # digit_img = digit_img.resize((digit_img.size[0] * 5, digit_img.size[1] * 5))

        digit_img_cv2 = numpy.array(digit_img)
        digit = pytesseract.image_to_string(digit_img_cv2,
            config='--psm 10 -c tessedit_char_whitelist=12345678')

        # print(digit)
        # digit_img.save(OUTPUT_IMAGE)
        # break

        try:
            digit = int(digit) % 10
            if 1 <= digit <= 8:
                prev_digit = digit
                draw.rectangle(((x_start, y_start), (x_end, y_end)), fill=colors[digit - 1])
            else:
                draw.rectangle(((x_start, y_start), (x_end, y_end)), fill=colors[prev_digit - 1])
        except:
            draw.rectangle(((x_start, y_start), (x_end, y_end)), fill=colors[prev_digit - 1])

img = img.filter(ImageFilter.BLUR)
img.save(OUTPUT_IMAGE)

# img = cv2.imread('Unpainted.png')
# print(pytesseract.image_to_string(img))