import numpy as np
from SIFT.GaussianFilterMatlab import gauss2D_matlab
from SIFT.DerotatePatch import derotate_patch
from scipy import ndimage


def compute_descriptors(blurred_images, keypoint_locations, rotation_invariant = 1):

    descriptors = {}
    final_kpt_locations = {}

    for octave in keypoint_locations.keys():
        # define gaussian filter with magic multiplication of 1.5 taken from Lowe's paper
        gaussian_filter = gauss2D_matlab((16, 16), 16 * 1.5)
        # check what images actually contain keypoints s.t. only relevant image indexes have to be searched
        relevant_img_idx = list(set(list[2] for list in keypoint_locations[octave]))
        # only gp through relevant images
        for img_idx in relevant_img_idx:

            # calculate image gradient magnitudes and directions
            grad_x, grad_y = np.gradient(blurred_images[octave][img_idx])
            grad_mag = np.sqrt(grad_x**2.0 + grad_y**2.0)
            grad_dir = np.arctan2(grad_y, grad_x) * (180 / np.pi)

            # get coordinates of keypoints for relevant image index
            img_kp_coord = [t[0:2] for t in keypoint_locations[octave] if t[2] == img_idx]

            # generate empty descriptor array
            img_descriptors = np.zeros([len(img_kp_coord), 128])

            # generate empty locator validation array - only keypoints far enough from image border are valid
            # reason: descriptor requires gradient information of surrounding pixels
            is_valid_locator = np.zeros(len(img_kp_coord))

            # get coordinates for each keypoint
            for kp_idx in range(len(img_kp_coord)):
                kp_row = img_kp_coord[kp_idx][0]
                kp_col = img_kp_coord[kp_idx][1]
                img_row_size, img_col_size = blurred_images[octave][img_idx].shape

                # TODO make sure that keypoints are not too close (not same spot); example - first keypoint
                if kp_row > 7 and kp_col > 7 and kp_row < img_row_size - 7 and kp_col < img_col_size - 7:
                    is_valid_locator[kp_idx] = 1
                    grad_mag_loc = grad_mag[kp_row - 8: kp_row + 8, kp_col - 8: kp_col + 8]
                    grad_mag_loc_weighted = grad_mag_loc * gaussian_filter
                    grad_dir_loc = grad_dir[kp_row - 8: kp_row + 8, kp_col - 8: kp_col + 8]

                    grad_mag_loc_derotated_weighted = grad_mag_loc_weighted
                    grad_dir_loc_derotated = grad_dir_loc

                    if rotation_invariant == 1:
                        # compute dominatnt direction through looking at the most common orientation in the histogram,
                        # spaced at 10 degrees --> from -180° to 180 results in 37 bins
                        bins_angles = np.arange(-180, 180, 10)
                        orientation_hist = np.histogram(grad_dir_loc, bins=bins_angles, weights=grad_mag_loc_weighted)
                        bin_dominant_direction = orientation_hist[1][orientation_hist[0].argmax()]
                        # take average of the two bin edges to calculate principal direction
                        grad_dir_loc_principal = (bin_dominant_direction + (bin_dominant_direction + 10)) / 2

                        # derotate patch
                        patch_derotated = derotate_patch(blurred_images[octave][img_idx], [kp_row, kp_col], 16, grad_dir_loc_principal)
                        grad_x, grad_y = np.gradient(patch_derotated)
                        grad_mag_loc_derotated = np.sqrt(grad_x ** 2.0 + grad_y ** 2.0)
                        grad_dir_loc_derotated = np.arctan2(grad_y, grad_x) * (180 / np.pi)
                        grad_mag_loc_derotated_weighted = grad_mag_loc_derotated * gaussian_filter

                    # Divide the 16×16 patch into 4 sub-patches of 4×4 size.
                    # For each sub-block, an 8 bin orientation histogram is created.
                    # Before creating the histogram, each gradient is weighted according to its scaled norm.
                    N_tmp = 0
                    bins_angles = np.arange(-180, 180, 45)
                    for ix in range(1, 4, 1):
                        for iy in range(1, 4, 1):
                            N_w = np.histogram(grad_dir_loc_derotated[4 * ix - 4: 4 * ix, 4 * iy - 4: 4 * iy],
                                               bins=bins_angles,
                                               weights=grad_mag_loc_derotated_weighted[4 * ix - 4: 4 * ix, 4 * iy - 4: 4 * iy])
                            img_descriptors[kp_idx, N_tmp: N_tmp + 7] = N_w[0][: 7]
                            N_tmp = N_tmp + 8;

            # Adapt keypoint location such that they correspond to the originial image dimensions.
            # TODO: The code crashes here
            img_kp_coord = [i * 2^(octave) for i in img_kp_coord]

            # Only store valid keypoints.
            for kp_idx in range(len(img_kp_coord)):
                if is_valid_locator[kp_idx] == 0:
                    img_descriptors = np.delete(img_descriptors, kp_idx, axis=0)
                    img_kp_coord = np.delete(img_kp_coord, kp_idx, axis=0)
            descriptors[octave] = img_descriptors
            final_kpt_locations[octave] = img_kp_coord

            # Normalize the descriptors such that they have unit norm.
            norm = np.linalg.norm(descriptors[octave])
            descriptors[octave] = descriptors[octave] / norm
            norm = np.linalg.norm(final_kpt_locations[octave])
            final_kpt_locations[octave] = final_kpt_locations[octave] / norm

    return descriptors, final_kpt_locations
