import cv2
import numpy as np

paper_size = [1240, 1754]

#function taken from
#https://www.pyimagesearch.com/2014/08/25/4-point-opencv-getperspective-transform-example/
def order_points(pts):
    rect = np.zeros((4, 2), dtype = "float32")

    s = pts.sum(axis = 1)
    rect[0] = pts[np.argmin(s)]
    rect[2] = pts[np.argmax(s)]

    diff = np.diff(pts, axis = 1)
    rect[1] = pts[np.argmin(diff)]
    rect[3] = pts[np.argmax(diff)]
    return rect

def qos_of_corners(contour):
    return cv2.contourArea(contour) > 3000000.0


def doc_rectification(orig_pts, new_pts, img):
    M = cv2.getPerspectiveTransform(orig_pts, new_pts)
    dst = cv2.warpPerspective(img,M,(paper_size[0],paper_size[1]))
    return dst

def doc_algorithm(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    ret, thresh = cv2.threshold(gray,0,255,cv2.THRESH_OTSU)
    kernel = np.ones((3,3), np.uint8)
    opening = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations = 2)

    sure_bg = cv2.dilate(opening,kernel,iterations=3)
    dist_transform = cv2.distanceTransform(opening, cv2.DIST_L2,5)
    ret, sure_fg = cv2.threshold(dist_transform, 0.7*dist_transform.max(), 255,0)
    sure_fg = np.uint8(sure_fg)
    unknown = cv2.subtract(sure_bg, sure_fg)
    ret, markers = cv2.connectedComponents(sure_fg)

    markers = markers + 1
    markers[unknown==255] = 0
    markers = cv2.watershed(img, markers)
    markers[markers == 1] = 0
    markers[markers != 0] = 1
    m = cv2.convertScaleAbs(markers)
    ret,thresh = cv2.threshold(m,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
    h,w = thresh.shape[:2]
    thresh[0,:] = 0
    thresh[:,0] = 0
    thresh[h-1,:] = 0
    thresh[:,w-1] = 0
    res = cv2.bitwise_and(img,img,mask=thresh)

    im2,contours,hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    cnt = contours[0]

    perimeter = cv2.arcLength(cnt,True)
    delta = 0.01
    while delta < 1.0:
        epsilon = delta*cv2.arcLength(cnt,True)
        approx = cv2.approxPolyDP(cnt,epsilon,True)
        delta = delta + 0.01
        if approx.shape[0] == 4:
            break
    cv2.drawContours(res,[approx],0,(0,255,0),20)

    pts2 = order_points(np.float32([[0,0],[0,paper_size[1]],[paper_size[0],paper_size[1]],[paper_size[0],0]]))
    pts1 = order_points(np.float32([x[0] for x in approx ]))
    dst = doc_rectification(pts1, pts2, img)
    return qos_of_corners(approx), dst
