from compareImgs import compareImgs
import cv2 as cv
from glob import glob
from compareImgs_hist import compareImgs_hist
from compareImgs_hist_rgb import compareImgs_hist_rgb
from SIFT import SIFT
from pprint import pprint
import os, sys

# the directory of the image database
database_dir = "image.orig"


def retrieval(choice="1"):
    if choice == "1":
        src_rgb = cv.imread("beach.jpg")
        # print("You choose: %s - beach\n" % choice)
    if choice == "2":
        src_rgb = cv.imread("building.jpg")
        # print("You choose: %s - building\n" % choice)
    if choice == "3":
        src_rgb = cv.imread("bus.jpg")
        # print("You choose: %s - bus\n" % choice)
    if choice == "4":
        src_rgb = cv.imread("dinosaur.jpg")
        # print("You choose: %s - dinosaur\n" % choice)
    if choice == "5":
        src_rgb = cv.imread("flower.jpg")
        # print("You choose: %s - flower\n" % choice)
    if choice == "6":
        src_rgb = cv.imread("horse.jpg")
        # print("You choose: %s - horse\n" % choice)
    if choice == "7":
        src_rgb = cv.imread("man.jpg")
        # print("You choose: %s - man\n" % choice)

    min_diff = 1e50

    # src_input = cv.imread("man.jpg")
    # cv.imshow("Input", src_input)

    # change the image to gray scale
    src_gray = cv.cvtColor(src_rgb, cv.COLOR_BGR2GRAY)

    # read image database
    database = sorted(glob(database_dir + "/*.jpg"))

    raw_results = []
    i = 0
    for img in database[0:150]:
        # read image
        img_rgb = cv.imread(img)
        # convert to gray scale
        img_gray = cv.cvtColor(img_rgb, cv.COLOR_BGR2GRAY)
        # compare the two images
        diff_gray = compareImgs(src_gray, img_gray)

        # compare the two images by histogram, uncomment the following line to use histogram
        diff_rgb = compareImgs_hist_rgb(src_rgb, img_rgb)

        sift_match = SIFT(src_rgb, img_rgb)
        raw_results.append([img, diff_gray, diff_rgb, sift_match])

        # if i > 3:
        #     break
        i += 1

    # normalize
    min_diff_gray = min([raw_result[1] for raw_result in raw_results])
    max_diff_gray = max([raw_result[1] for raw_result in raw_results])
    min_diff_rgb = min([raw_result[2] for raw_result in raw_results])
    max_diff_rgb = max([raw_result[2] for raw_result in raw_results])
    min_diff_sift_match = min([raw_result[3] for raw_result in raw_results])
    max_diff_sift_match = max([raw_result[3] for raw_result in raw_results])

    for i in range(len(raw_results)):
        raw_results[i] = [
            raw_results[i][0],
            (raw_results[i][1] - min_diff_gray) / (max_diff_gray - min_diff_gray),
            (raw_results[i][2] - min_diff_rgb) / (max_diff_rgb - min_diff_rgb),
            (raw_results[i][3] - min_diff_sift_match) / (max_diff_sift_match - min_diff_sift_match),
        ]

    weight_diff_gray = 0.7
    weight_diff_rgb = 0.3
    weight_sift_match = 0

    for raw_result in raw_results:
        img, diff_gray, diff_rgb, sift_match = raw_result
        test = sum(
            [
                #
                weight_diff_gray * diff_gray,
                weight_diff_rgb * diff_rgb,
                weight_sift_match * (1 - sift_match),
            ]
        )
        if test <= min_diff:
            min_diff = test
            result = img

    result_category = result.replace("image.orig/", "")[0]

    print("r:" + result + " rc:" + result_category)
    return result_category