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(rgb1, rgb2):
    # TODO: Fix CIE L*a*b difference calculation
    rgb1 = np.array([rgb1], dtype=np.uint8)
    rgb2 = np.array([rgb1], dtype=np.uint8)
    lab1 = cv2.cvtColor(rgb1, cv2.COLOR_RGB2Lab)
    lab2 = cv2.cvtColor(rgb1, cv2.COLOR_RGB2Lab)
    print(lab1)
    print(lab2)
    return 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)
    print(color_difference(np.array([1, 1, 0]), (np.array([0, 0, 0]))))
    white = np.array([99, 99, 99])
    blue = np.array([0x2A, 0xAB, 0xE1])
    black = np.array([0, 0, 0])
    height, width, channels = warped.shape
    print(round(height / 2))
    color_difference(blue, white)
    for x in range(width):
        pixel_color = warped[round(height / 2), x]
        detected = [pixel_color[2], pixel_color[1], pixel_color[0]]
        white_diff = color_difference(white, detected)
        blue_diff = color_difference(blue, detected)
        black_diff = color_difference(black, detected)
        print("\n")
        print(x)
        print(detected)
        print(white_diff)
        print(blue_diff)
        print(black_diff)
        if white_diff < blue_diff and white_diff < black_diff:
            print("white")
        elif blue_diff < white_diff and blue_diff < black_diff:
            print("blue")
            print(x)
        elif black_diff < blue_diff and black_diff < white_diff:
            print("black")
            print(x)