import os
import glob
import cv2
import numpy as np

SZ=20


def deskew(img):
    m = cv2.moments(img)
    if abs(m['mu02']) < 1e-2:
        return img.copy()
    skew = m['mu11'] / m['mu02']
    M = np.float32([[1, skew, -0.5 * SZ * skew], [0, 1, 0]])
    img = cv2.warpAffine(img, M, (SZ, SZ), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_LINEAR)
    return img


def get_hog():
    winSize = (20, 20)
    blockSize = (8, 8)
    blockStride = (4, 4)
    cellSize = (8, 8)
    nbins = 9
    derivAperture = 1
    winSigma = -1.
    histogramNormType = 0
    L2HysThreshold = 0.2
    gammaCorrection = 1
    nlevels = 64
    signedGradient = True

    hog = cv2.HOGDescriptor(winSize, blockSize, blockStride, cellSize, nbins, derivAperture, winSigma,
                            histogramNormType, L2HysThreshold, gammaCorrection, nlevels, signedGradient)

    return hog
    affine_flags = cv2.WARP_INVERSE_MAP | cv2.INTER_LINEAR


def get_contour_precedence(contour, cols):
    tolerance_factor = 10
    origin = cv2.boundingRect(contour)
    return ((origin[1] // tolerance_factor) * tolerance_factor) * cols + origin[0]


hog = get_hog();

model = cv2.ml.SVM_load('model.yml')
cv2.namedWindow("im", cv2.WINDOW_NORMAL)
img = cv2.imread('1.jpg')
x, y, w, h = cv2.selectROI("im", img, False, False)
cropped = img[y:y+h, x:x+w]
cropped = cv2.cvtColor(cropped, cv2.COLOR_BGR2GRAY)
test_samples = []

ret, thresh = cv2.threshold(cropped, 180, 255, cv2.THRESH_BINARY_INV)
canny = cv2.Canny(thresh, 150, 200, apertureSize=3)
contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

drawing = canny.copy()
rect = canny.copy()
drawing = cv2.cvtColor(drawing, cv2.COLOR_GRAY2BGR)
rect = cv2.cvtColor(rect, cv2.COLOR_GRAY2BGR)
cv2.drawContours(drawing, contours, -1, (0, 255, 0))
cv2.imshow("a", drawing)
cv2.waitKey(0)
contours.sort(key=lambda x:get_contour_precedence(x, thresh.shape[1]))
first_x, _, first_w, _ = cv2.boundingRect(contours[0])
last_x = first_x + first_w
batch_test_samples = []
for cnt in contours:
    area = cv2.contourArea(cnt)
    # print(cv2.contourArea(cnt))
    if area > 10:
        [x,y,w,h] = cv2.boundingRect(cnt)
        # print(h)
        if 5 < h < 25:
            roi = cropped[y:y+h, x:x+w]
            row, col = roi.shape[:2]
            bottom = roi[row - 2:row, 0:col]
            bordersize = 5
            border = cv2.copyMakeBorder(roi, top=bordersize, bottom=bordersize, left=bordersize, right=bordersize,
                                        borderType=cv2.BORDER_CONSTANT, value=[255,255,255])
            roiresize = cv2.resize(border, (SZ, SZ), interpolation=cv2.INTER_AREA)
            # cv2.imshow("roi", roiresize)
            # cv2.waitKey(0)
            if x - last_x < 10:
                batch_test_samples.append(roiresize)
            else:
                batch_hog_descriptors = []
                for img in np.array(batch_test_samples).reshape(-1, 20, 20):
                    batch_hog_descriptors.append(hog.compute(img))
                batch_hog_descriptors = np.squeeze(batch_hog_descriptors)
                test_samples.append(batch_hog_descriptors)
                batch_test_samples = [roiresize]
            last_x = x + w
            cv2.rectangle(rect, (x, y), (x + w, y + h), (0, 255, 0))

cv2.imshow("c", rect)
cv2.waitKey(0)
# resized = cv2.resize(cropped, (20, 20), interpolation=cv2.INTER_AREA)
# gray = cv2.cvtColor(resized, cv2.COLOR_BGR2GRAY)

# test_samples = np.array(test_samples)
# test_samples = test_samples.reshape(-1, SZ, SZ)

# print('Deskew images ... ')
# deskewed = list(map(deskew, test_samples))

# t1 = cv2.getTickCount()
# print('Defining HoG parameters ...')
# HoG feature descriptor

# print('Calculating HoG descriptor for every image ... ')
# hog_descriptors = []
# for img in deskewed:
#     hog_descriptors.append(hog.compute(img))
# hog_descriptors = np.squeeze(hog_descriptors)
# hog_descriptors = np.array(hog_descriptors)
for batch in test_samples:
    print(batch.shape)
    print(model.predict(batch)[1].ravel())

# t2 = cv2.getTickCount()
# print((t2 - t1)/cv2.getTickFrequency())