import cv2
import matplotlib.pyplot as plt
import torch
from PIL.Image import Image
from torch import nn
from torchvision.transforms import transforms
from torch.autograd import Variable
import numpy as np

from nn_model import Net


def recognizer(a_path):
    code = []
    path = a_path

    transform = transforms.Compose([transforms.ToTensor(),
                                    transforms.Normalize((0.5,), (0.5,)),
                                    ])

    img = cv2.imread(path)

    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img_gray = cv2.GaussianBlur(img_gray, (5, 5), 0)

    ret, im_th = cv2.threshold(img_gray, 90, 255, cv2.THRESH_BINARY_INV)

    ctrs, hier = cv2.findContours(im_th.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    rects = [cv2.boundingRect(ctr) for ctr in ctrs]

    # load nn model
    input_size = 784  # = 28*28
    hidden_sizes = [128, 128, 64]
    output_size = 10
    model = Net()
    model.load_state_dict(torch.load('model.pt'))
    model.eval()

    for rect in rects:
        # Crop image
        crop_img = img[rect[1]:rect[1] + rect[3] + 10, rect[0]:rect[0] + rect[2] + 10, 0]
        # Resize the image
        roi = cv2.resize(crop_img, (28, 28), interpolation=cv2.INTER_CUBIC)
        # roi = cv2.dilate(roi, (3, 3))
        # plt.imshow(roi)
        # plt.show()
        im = transform(roi)
        im = im.view(1, 1, 28, 28)
        with torch.no_grad():
            logps = model(im)
        ps = torch.exp(logps)
        probab = list(ps.numpy()[0])
        code.append(probab.index(max(probab)))

    print(code)
    # cv2.imshow("Code", img)
    # cv2.waitKey()
    return code


recognizer("55555.jpg")
# print(recognizer("55555.jpg"))