from __future__ import print_function
import sys
import cv2 as cv
import numpy as np
from hsvfilter import HsvFilter

use_mask = False
img = None
templ = None
mask = None
image_window = "Source Image"
result_window = "Result window"
match_method = 3
max_Trackbar = 5


def main():


    global img
    global templ
    img = cv.imread("equip/rings/test_screen.jpg", cv.IMREAD_COLOR)
    templ = cv.imread("equip/rings/ring_1_32.jpg", cv.IMREAD_COLOR)

    hsv = HsvFilter(13, 40, 85, 135, 255, 255, 0, 0, 55, 53)
    #img = apply_hsv_filter(img, hsv)
    #templ = apply_hsv_filter(templ, hsv)


    global use_mask
    use_mask = True
    global mask
    mask = cv.imread("equip/rings/ring_1_32-mask.png", cv.IMREAD_COLOR)


    cv.namedWindow(image_window, cv.WINDOW_AUTOSIZE)
    cv.namedWindow(result_window, cv.WINDOW_AUTOSIZE)

    trackbar_label = 'Method: \n 0: SQDIFF \n 1: SQDIFF NORMED \n 2: TM CCORR \n 3: TM CCORR NORMED \n 4: TM COEFF \n 5: TM COEFF NORMED'
    cv.createTrackbar(trackbar_label, image_window, match_method, max_Trackbar, MatchingMethod)

    MatchingMethod(match_method)

    cv.waitKey(0)
    return 0

def draw_rectangles(haystack_img, rectangles):
    # these colors are actually BGR
    line_color = (0, 255, 0)
    line_type = cv.LINE_4
    pic = None
    for (x, y, w, h) in rectangles:
        # determine the box positions
        top_left = (x, y)
        bottom_right = (x + w, y + h)
        # draw the box
        cv.rectangle(haystack_img, top_left, bottom_right, line_color, lineType=line_type)

        #pic = haystack_img[y:y + h, x:x + w]

    return haystack_img

def shift_channel(c, amount):
    if amount > 0:
        lim = 255 - amount
        c[c >= lim] = 255
        c[c < lim] += amount
    elif amount < 0:
        amount = -amount
        lim = amount
        c[c <= lim] = 0
        c[c > lim] -= amount
    return c

def apply_hsv_filter(original_image, hsv_filter):
    # convert image to HSV
    hsv = cv.cvtColor(original_image, cv.COLOR_BGR2HSV)


    # add/subtract saturation and value
    h, s, v = cv.split(hsv)
    s = shift_channel(s, hsv_filter.sAdd)
    s = shift_channel(s, -hsv_filter.sSub)
    v = shift_channel(v, hsv_filter.vAdd)
    v = shift_channel(v, -hsv_filter.vSub)
    hsv = cv.merge([h, s, v])

    # Set minimum and maximum HSV values to display
    lower = np.array([hsv_filter.hMin, hsv_filter.sMin, hsv_filter.vMin])
    upper = np.array([hsv_filter.hMax, hsv_filter.sMax, hsv_filter.vMax])
    # Apply the thresholds
    mask = cv.inRange(hsv, lower, upper)
    result = cv.bitwise_and(hsv, hsv, mask=mask)

    # convert back to BGR for imshow() to display it properly
    img = cv.cvtColor(result, cv.COLOR_HSV2BGR)

    return img

def MatchingMethod(param):
    global match_method
    match_method = param

    img_display = img.copy()

    method_accepts_mask = (cv.TM_SQDIFF == match_method or match_method == cv.TM_CCORR_NORMED)
    if (use_mask and method_accepts_mask):
        result = cv.matchTemplate(img, templ, match_method, None, mask)
    else:
        result = cv.matchTemplate(img, templ, match_method)
    #_minVal, _maxVal, minLoc, maxLoc = cv.minMaxLoc(result, None)
    cv.normalize(result, result, 0, 1, cv.NORM_MINMAX, -1)
    #_minVal, _maxVal, minLoc, maxLoc = cv.minMaxLoc(result, None)

    locations = np.where(result >= 0.91)
    locations = list(zip(*locations[::-1]))

    needle_w = templ.shape[1]
    needle_h = templ.shape[0]

    find_num = 20
    idx_1d = np.argpartition(result.flatten(), -find_num)[-find_num:]
    #new_res = result.flatten()[idx_1d]
    #new_res.append()
    idx_2d = np.unravel_index(idx_1d, result.shape)

    rectangles = []
    for i in range(0, len(idx_2d[0]), 1):
        y = int(idx_2d[0][i])
        x = int(idx_2d[1][i])
        rect = [x, y, needle_w, needle_h]
        # Add every box to the list twice in order to retain single (non-overlapping) boxes
        rectangles.append(rect)
        rectangles.append(rect)


    for loc in locations:
        rect = [int(loc[0]), int(loc[1]), needle_w, needle_h]
        # Add every box to the list twice in order to retain single (non-overlapping) boxes
        #rectangles.append(rect)
        #rectangles.append(rect)


    rectangles, weights = cv.groupRectangles(rectangles, groupThreshold=1, eps=0.5)
    keep_rects = []
    for rect in rectangles:
        w = rect[0]
        h = rect[1]
        x = rect[2] + w
        y = rect[3] + h
        screenshot_pos = img_display[h:y, w:x]  # (w, h, x+w, y+h)
        result2 = cv.matchTemplate(screenshot_pos, templ, 5)
        _minVal2, _maxVal2, minLoc2, maxLoc2 = cv.minMaxLoc(result2, None)
        if _maxVal2 >= 0.5:
            keep_rects.append(rect)
            print("matching error")






    res1 = draw_rectangles(img_display, keep_rects)
    res2 = draw_rectangles(img, keep_rects)

    cv.imshow(image_window, res1)
    cv.imshow(result_window, res2)



if __name__ == "__main__":
    main()