import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

import histogram_module
import dist_module


def rgb2gray(rgb):
    r, g, b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
    gray = 0.2989 * r + 0.5870 * g + 0.1140 * b

    return gray


# model_images - list of file names of model images
# query_images - list of file names of query images
#
# dist_type - string which specifies distance type:  'chi2', 'l2', 'intersect'
# hist_type - string which specifies histogram type:  'grayvalue', 'dxdy', 'rgb', 'rg'
#
# note: use functions 'get_dist_by_name', 'get_hist_by_name' and 'is_grayvalue_hist' to obtain
#       handles to distance and histogram functions, and to find out whether histogram function
#       expects grayvalue or color image

def find_best_match(model_images, query_images, dist_type, hist_type, num_bins):
    hist_isgray = histogram_module.is_grayvalue_hist(hist_type)

    model_hists = compute_histograms(model_images, hist_type, hist_isgray, num_bins)
    query_hists = compute_histograms(query_images, hist_type, hist_isgray, num_bins)

    D = np.zeros((len(model_images), len(query_images)))

    # compute distance for each couple of query - image
    for j, query in enumerate(query_hists):
        for i, model in enumerate(model_hists):
            D[i, j] = dist_module.get_dist_by_name(model, query, dist_type)

    best_match = []  # to save best matches

    # for each query , find best model
    for j in range(len(query_images)):
        query_matches = D[:, j]  # get query columns from matrix
        argmin = np.argmin(query_matches)  # get index with minimum distance
        best_match.append(argmin)  # save index for query

    best_match = np.array(best_match)  # array of best match for each query

    return best_match, D


def compute_histograms(image_list, hist_type, hist_isgray, num_bins):
    image_hist = []

    # Compute hisgoram for each image and add it at the bottom of image_hist

    # ... (your code here)
    for img in image_list:
        img_color = np.array(Image.open(img))

        # if hist is gray type we use gray image
        # othewise rgb image
        img_to_process = rgb2gray(img_color) if hist_isgray else img_color.astype('double')

        # We compute histogram for image
        hist = histogram_module.get_hist_by_name(img=img_to_process,
                                                 num_bins_gray=num_bins,
                                                 hist_name=hist_type
                                                 )
        image_hist.append(hist)

    return image_hist


# For each image file from 'query_images' find and visualize the 5 nearest images from 'model_image'.
#
# Note: use the previously implemented function 'find_best_match'
# Note: use subplot command to show all the images in the same Python figure, one row per query image

def show_neighbors(model_images, query_images, dist_type, hist_type, num_bins):
    plt.figure()

    num_nearest = 5  # show the top-5 neighbors

    # ... (your code here)

    _, D = find_best_match(model_images=model_images,
                           query_images=query_images,
                           dist_type=dist_type,
                           hist_type=hist_type,
                           num_bins=num_bins
                           )

    Q = len(query_images)
    pos = 0
    for j in range(Q):
        query_matches = D[:, j]
        best_args = np.argsort(query_matches)[:num_nearest]

        query_img = query_images[j]

        pos += 1
        plt.subplot(Q, 6, pos);
        plt.imshow(np.array(Image.open(query_img)), vmin=0, vmax=255);
        plt.title(f'Q{j}')
        for ind in range(len(best_args)):
            pos += 1
            model_ind = best_args[ind]
            model_img = model_images[model_ind]
            plt.subplot(Q, 6, pos);
            plt.imshow(np.array(Image.open(model_img)), vmin=0, vmax=255);
            plt.title(f'MO.{model_ind}')

    plt.show()




[((27, 18, 27, 63), 'histogram_module.is_grayvalue_hist', 'histogram_module.is_grayvalue_hist', ({(27, 53, 27, 62): 'hist_type'}, {}), '(hist_type)', False, 'import histogram_module\n'), ((47, 17, 47, 37), 'numpy.array', 'np.array', ({(47, 26, 47, 36): 'best_match'}, {}), '(best_match)', True, 'import numpy as np\n'), ((81, 4, 81, 16), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((114, 4, 114, 14), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((44, 17, 44, 41), 'numpy.argmin', 'np.argmin', ({(44, 27, 44, 40): 'query_matches'}, {}), '(query_matches)', True, 'import numpy as np\n'), ((66, 15, 69, 50), 'histogram_module.get_hist_by_name', 'histogram_module.get_hist_by_name', (), '', False, 'import histogram_module\n'), ((103, 8, 103, 30), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(103, 20, 103, 21): 'Q', (103, 23, 103, 24): '(6)', (103, 26, 103, 29): 'pos'}, {}), '(Q, 6, pos)', True, 'import matplotlib.pyplot as plt\n'), ((105, 8, 105, 26), 'matplotlib.pyplot.title', 'plt.title', ({(105, 18, 105, 25): 'f"""Q{j}"""'}, {}), "(f'Q{j}')", True, 'import matplotlib.pyplot as plt\n'), ((37, 22, 37, 75), 'dist_module.get_dist_by_name', 'dist_module.get_dist_by_name', ({(37, 51, 37, 56): 'model', (37, 58, 37, 63): 'query', (37, 65, 37, 74): 'dist_type'}, {}), '(model, query, dist_type)', False, 'import dist_module\n'), ((59, 29, 59, 44), 'PIL.Image.open', 'Image.open', ({(59, 40, 59, 43): 'img'}, {}), '(img)', False, 'from PIL import Image\n'), ((98, 20, 98, 45), 'numpy.argsort', 'np.argsort', ({(98, 31, 98, 44): 'query_matches'}, {}), '(query_matches)', True, 'import numpy as np\n'), ((110, 12, 110, 34), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(110, 24, 110, 25): 'Q', (110, 27, 110, 28): '(6)', (110, 30, 110, 33): 'pos'}, {}), '(Q, 6, pos)', True, 'import matplotlib.pyplot as plt\n'), ((112, 12, 112, 40), 'matplotlib.pyplot.title', 'plt.title', ({(112, 22, 112, 39): 'f"""MO.{model_ind}"""'}, {}), "(f'MO.{model_ind}')", True, 'import matplotlib.pyplot as plt\n'), ((104, 28, 104, 49), 'PIL.Image.open', 'Image.open', ({(104, 39, 104, 48): 'query_img'}, {}), '(query_img)', False, 'from PIL import Image\n'), ((111, 32, 111, 53), 'PIL.Image.open', 'Image.open', ({(111, 43, 111, 52): 'model_img'}, {}), '(model_img)', False, 'from PIL import Image\n')]