# -*- coding: utf-8 -*-
# 参考 https://www.jianshu.com/p/2bbdb27ee7b3 
#python3 t1.py --image images/IMG20190622074037.jpg

# 引入必要的库
from imutils.perspective import four_point_transform
from imutils import contours
import numpy as np
import argparse
import imutils
import cv2
import time

def showimg(img,name='cusom image'):
    # cv2.namedWindow(name, cv2.WINDOW_NORMAL)
    cv2.imshow(name, img)
    cv2.waitKey(0)
    # cv2.destroyAllWindows()

# 构建命令行参数解析并分析参数
# 对应使用方式 python3 t1.py --image images/IMG20190622074037.jpg
ap = argparse.ArgumentParser()
ap.add_argument("-i", "--image", required=True,
    help="path to the input image")
args = vars(ap.parse_args())

# 构建答案字典，键为题目号，值为正确答案
ANSWER_KEY = {0: 1, 1: 4, 2: 0, 3: 3, 4: 1}

# 加载图片，将它转换为灰阶，轻度模糊，然后边缘检测。
image = cv2.imread(args["image"])
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
blurred = cv2.GaussianBlur(gray, (5, 5), 0)
edged = cv2.Canny(blurred, 75, 400)

## 展示图片
# showimg(edged)

start_time = time.time()

# 从边缘图中寻找轮廓，然后初始化答题卡对应的轮廓
cnts = cv2.findContours(edged.copy(), cv2.RETR_EXTERNAL,
    cv2.CHAIN_APPROX_SIMPLE)
if imutils.is_cv2() or imutils.is_cv4():
    cnts = cnts[0]
else:
    cnts = cnts[1]

docCnt = None
 
# 确保至少有一个轮廓被找到
if len(cnts) > 0:
    # 将轮廓按大小降序排序
    cnts = sorted(cnts, key=cv2.contourArea, reverse=True)
 
    # 对排序后的轮廓循环处理
    for c in cnts:
        # 获取近似的轮廓
        peri = cv2.arcLength(c, True)
        approx = cv2.approxPolyDP(c, 0.02 * peri, True)
 
        # 如果我们的近似轮廓有四个顶点，那么就认为找到了答题卡
        if len(approx) == 4:
            docCnt = approx
            break

# print(docCnt)  #得到4个坐标

# 对原始图像和灰度图都进行四点透视变换
paper = four_point_transform(image, docCnt.reshape(4, 2))
warped = four_point_transform(gray, docCnt.reshape(4, 2))
# warped2 = four_point_transform(image,    np.array([ [79,81],[29,1035],[742,1044],[734,99]  ]  )       )
# warped3 = four_point_transform(image,    np.array([ [79,111],[59,1035],[742,1004],[734,99]  ]  )       )

# # 展示图片
# showimg(warped2)
# showimg(warped3)


############## 至此得到透视变换后的图像 ###########

############## 接下来开始做：以科级蓝表p2为例          ###########
'''
1.把透视变换后的图像处理为宽800高自动的图像，以便接下来基于间距定位进行处理。
2.裁剪图像边缘-20像素，以避免待会查找直线时查找到边缘线
3.查找直线，找出靠右的5条直线和靠下的9条直线，得到所有交点，  # 二期：如果超出范围就改变参数再次查找
4.做一次透视变换得到比较方正的评分区域
5.再次裁剪边缘-10像素，以便排除边缘干扰
6.重新查找直线，获取交点坐标（这样就可以免去坐标转换的麻烦），记住交点坐标
7.去除直线（避免干扰填写判断）
8.对交点坐标划分出的单元格进行权重运算，高权重者即为被选项

【方案二】
1.把透视变换后的图像处理为宽800高自动的图像，以便接下来基于间距定位进行处理。
2.裁剪图像边缘-20像素，以避免待会查找直线时查找到边缘线
3.查找直线，得到所有交点（调大膨胀参数以便得到尽可能多的交点，之后再通过算法合并过于靠近的点）
4.对生成交点的图像进行膨胀，以便清除太相近的点，最后提取这些‘交点坐标’
5.算法计算相应单元格对应坐标（注意有偏移）
6.对交点坐标划分出的单元格进行透视变换/去除直线/裁剪等组合操作去除单元格边框干扰
7.对交点坐标划分出的单元格进行权重运算，高权重者即为被选项
'''

# 1.把透视变换后的图像处理为宽800高自动的图像，以便接下来基于间距定位进行处理。

x, y = warped.shape[0:2]
width = 800
r = width/y
dim = (width, int(x*r))
img800 =   cv2.resize(warped, dim, interpolation=cv2.INTER_AREA)
#### cv2.imwrite("./images_save/img800.jpg",img800) 
#showimg(img800)


# 2.裁剪图像边缘-20像素，以避免待会查找直线时查找到边缘线

x , y = img800.shape[0:2]
indent = 15 #裁剪像素
cropped = img800[indent:x-indent, indent:y-indent] #startX:endX, startY:endY
#### cv2.imwrite("./images_save/cropped.jpg",cropped)


# 3.查找直线，找出靠右的5条直线和靠下的9条直线，得到所有交点，  # 二期：如果超出范围就改变参数再次查找

#二值化
gray = cropped
binary = cv2.adaptiveThreshold(~gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 15, -10)
# cv2.imshow("cell", binary)
#### cv2.imwrite("./images_save/binary.jpg",binary)
# cv2.waitKey(0)
rows,cols=binary.shape
scale = 50 #关系到识别长短
erode = 2 #腐蚀参数 
dilate = 5 #膨胀参数
#识别横线
kernel  = cv2.getStructuringElement(cv2.MORPH_RECT,(cols//scale,1))
eroded = cv2.erode(binary,kernel,iterations = erode)
#cv2.imshow("Eroded Image",eroded)
dilatedcol = cv2.dilate(eroded,kernel,iterations = dilate)
#cv2.imshow("Dilated Image",dilatedcol)
#### cv2.imwrite("./images_save/dilatedcol.jpg",dilatedcol)
cv2.waitKey(0)
#识别竖线
kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(1,rows//scale))
eroded = cv2.erode(binary,kernel,iterations = erode)
dilatedrow = cv2.dilate(eroded,kernel,iterations = dilate)
#cv2.imshow("Dilated Image",dilatedrow)
#### cv2.imwrite("./images_save/dilatedrow.jpg",dilatedrow)
cv2.waitKey(0)
# 标识交点
bitwiseAnd = cv2.bitwise_and(dilatedcol,dilatedrow)
# cv2.imshow("bitwiseAnd Image",bitwiseAnd)
#### cv2.imwrite("./images_save/bitwiseAnd.jpg",bitwiseAnd)
cv2.waitKey(0)
#标识表格
# merge = cv2.add(dilatedcol,dilatedrow)
# cv2.imshow("add Image",merge)
# cv2.imwrite("./images_save/merge.jpg",merge)
# cv2.waitKey(0)


# 4.对生成交点的图像进行膨胀，以便清除太相近的点，最后提取这些‘交点坐标’

kernel  = cv2.getStructuringElement(cv2.MORPH_RECT,(4,4))
dilateddot = cv2.dilate(bitwiseAnd,kernel,iterations = 1)
# cv2.imshow("dilateddot Image 2",dilateddot)
cv2.waitKey(0)
#### cv2.imwrite("./images_save/dilateddot.jpg",dilateddot)
# 提取坐标
dots = []
cnts = cv2.findContours(dilateddot, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if imutils.is_cv2() or imutils.is_cv4():
    cnts = cnts[0]
else:
    cnts = cnts[1]
for c in cnts:
    (x, y), radius = cv2.minEnclosingCircle(c)# 计算各个轮廓的中心
    # print(x,y)
    center = (int(x), int(y))# 转为整数 cast to integers
    # print(center)
    dots.append(center)
    radius = int(radius)
    dilateddot2 = cv2.circle(dilateddot, center, radius, (255, 255, 255), 2)# 绘圆 draw the circle
    cropped2 = cv2.circle(cropped, center, 2, (0, 0, 0), 2)# 绘圆 draw the circle

#### cv2.imwrite("./images_save/dilateddot2.jpg",dilateddot2)
#### cv2.imwrite("./images_save/cropped_2.jpg",cropped2)
# cv2.imshow("dilateddot2 Image 2",cropped2)
cv2.waitKey(0)

# print(dots)  


########  至此得到了各交点坐标。 但是请注意，分割的单元格四点，坐标存在偏移



# 5.算法计算相应单元格对应坐标（注意有偏移）  ##########

'''
  这个过程稍微复杂，而且依赖于表格特性
  首先我们只需要表格右边的选项填写区域，这些区域单元格有一个最大宽度max_cell_width
  以及最小表格高度min_cell_height。
  还有就是横线的坐标纵向最大偏移max_hori_jump用于确定下一行的单元格坐标，
  本程序是固定判断某页面所以固定为查找到40个格子就结束
'''
max_cell_width  = 60
min_cell_height = 30
max_hori_jump   = 15
xColCells       = 4   # 选项区一行有4列（4个单元格）
xAllCells       = 40  # 一共有40个单元格


# 注意这些坐标【基本上】是按纸上的位置从右下角为起点，向左再向上的顺序排列的，恰好符合我们提取格子的顺序
#      但是，调试过程才发现，纵向是规律的从下往上，但是横向的却不总是从右往左，存在跳跃。因此，要做处理

# 5-1 对交点坐标进行梳理
#      梳理的思路是首先区分出横向线集合(根据max_hori_jump区分)，再对集合内进行排序
dots2 = []
horiOneSet = [] # 横向线集合
dots_len = len(dots)
i = 0
while i < dots_len: # 0-
    horiOneSet.append(  (dots[i][0],dots[i][1])  )
    if i == dots_len-1: #到了最后一个
        #排序已压入的,并清空
        horiOneSet.sort()
        horiOneSet.reverse()
        dots2.extend(horiOneSet)
        horiOneSet = []
    elif i>0 and (abs( dots[i+1][1] - dots[i][1] ) > max_hori_jump) :  # 如果下一个跳跃
        #排序已压入的,并清空
        horiOneSet.sort()
        horiOneSet.reverse()
        dots2.extend(horiOneSet)
        horiOneSet = []
    i += 1

dots = dots2

# 5-2 查找并绘制单元格矩形
xCell    = []   # 查找结果 这里面要放坐标各个单元格的顶点坐标 如[[A,B,C,D]]如A是[x,y]
i = 0
dots_len = len(dots)
while i < dots_len:
    oneA = (dots[i][0],   dots[i][1])    # d[i].x , d[i].y
    oneB = (dots[i+1][0],   dots[i+1][1]) # d[i+1].x , d[i+1].y
    j = i + 1
    m = 0
    while j < dots_len: # 查找下一行，下一行应偏移大于最小单元格高而且不在直线误差范围内
        if ( abs(dots[j][1] - dots[i][1]) > min_cell_height) and (abs(dots[j][1]-dots[j-1][1])>max_hori_jump) :
            m = j+len(xCell)%xColCells
            break
        j += 1
    if m==0 :
        print('ERROR m==0')
    oneC = ( dots[m][0], dots[m][1]  )
    oneD = ( dots[m+1][0], dots[m+1][1]  )

    one = (oneA,oneB,oneD,oneC)
    xCell.append(one)
    if len(xCell)<=xAllCells:
        # print(one)
        pts = np.array(one, np.int32)
        pts = pts.reshape((-1,1,2))
        cv2.polylines(cropped2,[pts],True,(0,255,255)) # ####

    if len(xCell) == xAllCells :
        break
    # 够4列换行
    if len(xCell)%xColCells==0 :  #利用最大表格宽度筛选也可以 abs(dots[i+1][0] - dots[i][0])  > max_cell_width 
        i = m - 3
    else:
        i += 1

# print(xCell)


#### cv2.imwrite("./images_save/cropped_2_polylines.jpg",cropped2)
# cv2.imshow("dilateddot2 Image 2",cropped2)
# cv2.waitKey(0)




# 6.对交点坐标划分出的单元格进行透视变换/去除直线/裁剪等组合操作去除单元格边框干扰
############################################################################

# 注意xCell[i]是ABDC的形式 A右下角B左下角D左上角C右上角 
#  D      C
#
#  B      A
# xCell reverse之后就是从上到下从左到右的顺序了     )
xCell.reverse()
# print('xCell.len',len(xCell))
i = 0 
points = []  # 每个单元的权重积分（这里采用的是计算白色像素个数）
checked= []  # 保存目标选中项，每行保存一个索引值。例如4x10的填写区，只会生成10个记录，每个记录保存选项对应的索引。
colorImg = cv2.cvtColor(cropped2, cv2.COLOR_GRAY2RGB)
for cell in xCell:
    (A,B,D,C) = cell
    # print(A,B,D,C)
    # 透视变换
    cell_toushi = four_point_transform( cropped2,    np.array([ D,B,A,C ]  ) )
    # 切出单元格图片边缘
    x , y = cell_toushi.shape[0:2]
    indent = 4 #像素
    cell_toushi_cut = cell_toushi[indent:x-indent, indent:y-indent] #startX:endX, startY:endY
    # cv2.imwrite("./images_save/cell_toushi_cut_%s.jpg" % (i),cell_toushi_cut)

    #比较  二值化后计算白色像素个数
    gray = cell_toushi_cut
    binary = cv2.adaptiveThreshold(~gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 15, -10)
    height, width = binary.shape
    area = 0
    for m in range(height):
        for n in range(width):
            if binary[m, n] == 255:
                area += 1
    # print (i,area)
    points.append(area)
    if (i+1)%4==0 :
        maxi = max( [points[i-3],points[i-2],points[i-1],points[i]]   )
        # print('maxi',maxi)
        rowChecks = [points[i-3],points[i-2],points[i-1],points[i]]
        rowCheckedIdx = rowChecks.index(maxi) 
        # print('rowChecks',rowChecks)
        # print('rowCheckedIdx',rowCheckedIdx)
        checked.append( rowCheckedIdx ) 
        # draw a green line around on xCell[i-4+rowCheckedIdx]
        # print('i-4+rowCheckedIdx',i-3+rowCheckedIdx)
        pts = np.array( (xCell[i-3+rowCheckedIdx]), np.int32)
        pts = pts.reshape((-1,1,2))
        cv2.polylines(colorImg,[pts],True,(0,255,0)) # ####

    # cv2.imwrite("./images_save/_cbinary_%s(%s).jpg" % (i,area),binary)

    i +=1


cv2.imwrite("./images_save/_colorImg.jpg" ,colorImg)
print('the answer is',checked)
end_time = time.time()
print("durations %s s" % (end_time-start_time) )