import cv2
import os
import numpy as np
from sklearn import svm
from sklearn import neighbors
from sklearn.neural_network import MLPClassifier


# 垂直投影，分割字符
def verticle_projection(thresh1):
    h, w = thresh1.shape[:]
    a = [0] * w
    for j in range(w):
        for i in range(h):
            if thresh1[i, j] == 0:
                a[j] += 1
                thresh1[i, j] = 255
    
    for j in range(w):
        for i in range(h - a[j], h):
            thresh1[i, j] = 0
    roi_list = []
    start_index = 0
    end_index = 0
    in_block = False
    for i in range(w):
        if in_block == False and a[i] != 0:
            in_block = True
            start_index = i
        elif a[i] == 1 and in_block:
            end_index = i
            in_block = False
        roiImg = thresh1[:h, start_index: end_index + 1]
        roi_list.append(roiImg)
    return roi_list

# 提取网格特征
def get_features(array):
    h, w = array.shape[:]
    data = []
    for x in range(w // 4):
        offset_y = x * 4
        temp = []
        for y in range(h // 4):
            offset_x = y * 4
            sum_temp = array[0 + offset_y: 4 + offset_y, 0 + offset_x: 4 + offset_x]
            temp.append(sum(sum(sum_temp)))
        data.append(temp)
    return np.asarray(data)


def train():
    train_path = './dataset/train/'
    train_files = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
    train_X = []
    train_Y = []
    for train_file in train_files:
        pictures = os.listdir(train_path + train_file)
        for picture in pictures:
            img = cv2.imread(train_path + train_file + '/' + picture)
            img = cv2.resize(img, (32, 32), cv2.INTER_CUBIC)
            gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            ret, threshold = cv2.threshold(gray_img, 130, 255, cv2.THRESH_BINARY)
            feature = get_features(threshold)
            feature = feature.reshape(feature.shape[0] * feature.shape[1])
            train_X.append(feature)
            train_Y.append(train_file)
    train_X = np.array(train_X)
    train_Y = np.array(train_Y)
    # svm
    linearsvc_clf = svm.LinearSVC()
    linearsvc_clf.fit(train_X, train_Y)
    # knn
    knn_clf = neighbors.KNeighborsClassifier()
    knn_clf.fit(train_X, train_Y)
    # mlp
    mlp_clf = MLPClassifier(solver='lbfgs', alpha=1e-5, hidden_layer_sizes=(32, 32), random_state=1)
    mlp_clf.fit(train_X, train_Y)
    return linearsvc_clf


def test_main(linearsvc_clf):
    img = cv2.imread('')
    gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    ret, threshold = cv2.threshold(gray_img, 130, 255, cv2.THRESH_BINARY)
    roi_list = verticle_projection(threshold)  
    test_X = []
    for single in roi_list:
        single = cv2.resize(single, (32, 32),  cv2.INTER_CUBIC)
        feature = get_features(single)
        feature = feature.reshape(feature.shape[0] * feature.shape[1])
        test_X.append(feature)
    test_X = np.asarray(test_X)
    result = linearsvc_clf.predict(test_X)
    return result


if __name__ == '__main__':
    linearsvc_clf = train()
    test_main(linearsvc_clf)