import cv2
import imutils
import numpy as np


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 four_point_transform(image, pts):
    rect = order_points(pts)
    (tl, tr, br, bl) = rect
    width_a = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
    width_b = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
    max_width = max(int(width_a), int(width_b))

    height_a = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
    height_b = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
    max_height = max(int(height_a), int(height_b))

    dst = np.array([
        [0, 0],
        [max_width - 1, 0],
        [max_width - 1, max_height - 1],
        [0, max_height - 1]], dtype="float32")

    transformation = cv2.getPerspectiveTransform(rect, dst)
    warped = cv2.warpPerspective(image, transformation, (max_width, max_height))

    return warped


def align_images(im1, im2):
    im1_gra = cv2.cvtColor(im1, cv2.COLOR_BGR2GRAY)
    im2_gray = cv2.cvtColor(im2, cv2.COLOR_BGR2GRAY)

    orb = cv2.ORB_create(500)
    keypoints1, descriptors1 = orb.detectAndCompute(im1_gra, None)
    keypoints2, descriptors2 = orb.detectAndCompute(im2_gray, None)

    matcher = cv2.DescriptorMatcher_create(cv2.DESCRIPTOR_MATCHER_BRUTEFORCE_HAMMING)
    matches = matcher.match(descriptors1, descriptors2, None)

    matches.sort(key=lambda x: x.distance, reverse=False)

    num_good_matches = int(len(matches) * 0.15)
    matches = matches[:num_good_matches]

    im_matches = cv2.drawMatches(im1, keypoints1, im2, keypoints2, matches, None)
    cv2.imwrite("matches.jpg", im_matches)

    points1 = np.zeros((len(matches), 2), dtype=np.float32)
    points2 = np.zeros((len(matches), 2), dtype=np.float32)

    for i, match in enumerate(matches):
        points1[i, :] = keypoints1[match.queryIdx].pt
        points2[i, :] = keypoints2[match.trainIdx].pt

    h, mask = cv2.findHomography(points1, points2, cv2.RANSAC)

    height, width, channels = im2.shape
    im1_reg = cv2.warpPerspective(im1, h, (width, height))

    return im1_reg, h


def color_difference(color1, color2):
    percentage1 = (color1[0] / 255 * 100 + color1[1] / 255 * 100 + color1[2] / 255 * 100) / 3
    percentage2 = (color2[0] / 255 * 100 + color2[1] / 255 * 100 + color2[2] / 255 * 100) / 3
    return percentage1 - percentage2 if percentage1 - percentage2 > 0 else 100


if __name__ == '__main__':
    image = cv2.imread("example.jpg")
    ratio = image.shape[0] / 500.0
    orig = image.copy()

    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    gray = cv2.GaussianBlur(gray, (5, 5), 0)
    edged = cv2.Canny(gray, 75, 200)

    cnts = cv2.findContours(edged.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)
    cnts = sorted(cnts, key=cv2.contourArea, reverse=True)[:5]

    for c in cnts:
        peri = cv2.arcLength(c, True)
        approx = cv2.approxPolyDP(c, 0.02 * peri, True)
        if len(approx) == 4:
            screenCnt = approx
            break

    cv2.drawContours(image, [screenCnt], -1, (0, 255, 0), 2)

    warped = four_point_transform(orig, screenCnt.reshape(4, 2) * ratio)
    cv2.imwrite("before.jpg", warped)

    ref_filename = "mask.jpg"
    im_reference = cv2.imread(ref_filename, cv2.IMREAD_COLOR)
    warped, h = align_images(warped, im_reference)
    print("Estimated homography : \n", h)
    out_filename = "aligned.jpg"
    cv2.imwrite(out_filename, warped)

    # cv2.imshow("Warped", warped)
    # cv2.waitKey(0)
    color = warped[392, 54]
    blue = [0x2A, 0xAB, 0xE1]
    black = [0, 0, 0]
    detected = [color[2], color[1], color[0]]
    if color_difference(blue, detected) < 30:
        print("Color is blue!")
    elif color_difference(black, detected) < 30:
        print("Color is black!")