import matplotlib.pyplot as plt
import numpy as np
import cv2
from skimage import measure
from sabhi_utils.image_utils.utils import load_image

import os
from skimage.filters.rank import entropy
from skimage.morphology import square
import copy
import time


class BlurDetector:

    def __init__(
        self,
        preprocessor=None,
        thresh=20
    ):

        self._preprocessor = preprocessor
        self._thresh = thresh

    def _get_image_center(self):

        return (int(self._w / 2.0), int(self._h / 2.0))

    def _visualize(self):
        # check to see if we are visualizing our output
        # compute the magnitude spectrum of the transform
        # display the original input image
        (fig, ax) = plt.subplots(1, 2, )
        ax[0].imshow(self._image, cmap="gray")
        ax[0].set_title("Input")
        ax[0].set_xticks([])
        ax[0].set_yticks([])
        # display the magnitude image
        ax[1].imshow(self._magnitude, cmap="gray")
        ax[1].set_title("Magnitude Spectrum")
        ax[1].set_xticks([])
        ax[1].set_yticks([])
        # show our plots
        plt.show()

    def _zero_center_inverse(self):
        fft_shift = self._fft_shift.copy()

        fft_shift[
            self._cY - self._size:self._cY + self._size,
            self._cX - self._size:self._cX + self._size
        ] = 0

        fft_shift = np.fft.ifftshift(fft_shift)
        recon = np.fft.ifft2(fft_shift)
        return recon

    def __call__(
        self,
        image,
        vis=False,
        size_ratio=8,
        mag_multi = 25
    ):

        image = load_image(image)
        self._image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        self._h, self._w = self._image.shape
        self._cX, self._cY = self._get_image_center()

        self._size = int(self._h/size_ratio)

        self._fft = np.fft.fft2(self._image)
        self._fft_shift = np.fft.fftshift(self._fft)

        self._recon = self._zero_center_inverse()

        self._magnitude = mag_multi * np.log(
            np.abs(self._recon)
        )

        self._mean = np.mean(self._magnitude)

        if vis:
            self._visualize()

        # the image will be considered "blurry" if the mean value of the
        # magnitudes is less than the threshold value
        # return (self._mean <= self._thresh, self._mean)
        return {
            'isDetected': self._mean <= self._thresh,
            'mean': self._mean
        }


class BrightspotDetector:
    def __init__(
        self,
        preprocess=[]
    ):

        self._preprocess = preprocess

    def __call__(
        self,
        image,
        percentage_limit=0.02,
        debug=False
    ):

        image = load_image(image)
        self._image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        self._blurred = cv2.GaussianBlur(self._image, (11, 11), 0)
        self._size = self._image.shape[0] * self._image.shape[1]
        self._percentage_limit = percentage_limit
        self._blob_size_limit = int(self._size * self._percentage_limit)

        # threshold the image to reveal light regions in the
        # blurred image

        self._thresh = cv2.threshold(
            self._blurred, 238, 255, cv2.THRESH_BINARY)[1]

        self._thresh = cv2.erode(self._thresh, None, iterations=2)
        self._thresh = cv2.dilate(self._thresh, None, iterations=4)

        # perform a connected component analysis on the thresholded
        # image, then initialize a mask to store only the "large"
        # components
        labels = measure.label(self._thresh, background=0)
        if debug:
            cv2.imshow("self._thresh", self._thresh)
            cv2.waitKey(0)
            print(labels)
        #mask = np.zeros(self._thresh.shape, dtype="uint8")
        # loop over the unique components
        for label in np.unique(labels):
            # if this is the background label, ignore it
            if label == 0:
                continue
            # otherwise, construct the label mask and count the
            # number of pixels
            labelMask = np.zeros(self._thresh.shape, dtype="uint8")
            labelMask[labels == label] = 255
            numPixels = cv2.countNonZero(labelMask)
            # if the number of pixels in the component is sufficiently
            # large, then add it to our mask of "large blobs"
            if numPixels > self._blob_size_limit:
                return {
                    'isDetected': True
                }

        # return (False, "Insignificant bright spots")
        return {
            'isDetected': False
        }


class BlurDetectorSpat(object):
    def __init__(self, downsampling_factor=4, num_scales=4, scale_start=3, entropy_filt_kernel_sze=7, sigma_s_RF_filter=15, sigma_r_RF_filter=0.25, num_iterations_RF_filter=3):
        self.downsampling_factor = downsampling_factor
        self.num_scales = num_scales
        self.scale_start = scale_start
        self.entropy_filt_kernel_sze = entropy_filt_kernel_sze
        self.sigma_s_RF_filter = sigma_s_RF_filter
        self.sigma_r_RF_filter = sigma_r_RF_filter
        self.num_iterations_RF_filter = num_iterations_RF_filter
        self.scales = self.createScalePyramid()
        self.__freqBands = []
        self.__dct_matrices = []
        self.freq_index = []

    def disp_progress(self, i, rows, old_progress):
        progress_dict = {10: '[|                  ] 10%',
                         20: '[| |                ] 20%',
                         30: '[| | |              ] 30%',
                         40: '[| | | |            ] 40%',
                         50: '[| | | | |          ] 50%',
                         60: '[| | | | | |        ] 60%',
                         70: '[| | | | | | |      ] 70%',
                         80: '[| | | | | | | |    ] 80%',
                         90: '[| | | | | | | | |  ] 90%',
                         100: '[| | | | | | | | | |] 100%'}

        i_done = i / rows * 100
        p_done = round(i_done / 10) * 10
        if(p_done != old_progress):
            os.system('cls' if os.name == 'nt' else 'clear')
            print(progress_dict[p_done])
            old_progress = p_done
        return(p_done)

    def createScalePyramid(self):
        scales = []
        for i in range(self.num_scales):
            # Scales would be 7, 15, 31, 63 ...
            scales.append((2**(self.scale_start + i)) - 1)
        return(scales)

    def computeImageGradientMagnitude(self, img):
        # Find x and y gradients
        __sobelx = cv2.Sobel(img, cv2.CV_64F, 1, 0,
                             borderType=cv2.BORDER_REFLECT)
        __sobely = cv2.Sobel(img, cv2.CV_64F, 0, 1,
                             borderType=cv2.BORDER_REFLECT)

        # Find gradient magnitude
        __magnitude = np.sqrt(__sobelx ** 2.0 + __sobely ** 2.0)
        return(__magnitude)

    def __computeFrequencyBands(self):
        for current_scale in self.scales:
            matrixInds = np.zeros((current_scale, current_scale))

            for i in range(current_scale):
                matrixInds[0: max(
                    0, int(((current_scale-1)/2) - i + 1)), i] = 1

            for i in range(current_scale):
                if (current_scale-((current_scale-1)/2) - i) <= 0:
                    matrixInds[0:current_scale - i - 1, i] = 2
                else:
                    matrixInds[int(current_scale - ((current_scale - 1) / 2) -
                                   i - 1): int(current_scale - i - 1), i] = 2
            matrixInds[0, 0] = 3
            self.__freqBands.append(matrixInds)

    def __dctmtx(self, n):
        [mesh_cols, mesh_rows] = np.meshgrid(
            np.linspace(0, n-1, n), np.linspace(0, n-1, n))
        dct_matrix = np.sqrt(
            2/n) * np.cos(np.pi * np.multiply((2 * mesh_cols + 1), mesh_rows) / (2*n))
        dct_matrix[0, :] = dct_matrix[0, :] / np.sqrt(2)
        return(dct_matrix)

    def __createDCT_Matrices(self):
        if(len(self.__dct_matrices) > 0):
            raise TypeError(
                "dct matrices are already defined. Redefinition is not allowed.")
        for curr_scale in self.scales:
            dct_matrix = self.__dctmtx(curr_scale)
            self.__dct_matrices.append(dct_matrix)

    def __getDCTCoefficients(self, img_blk, ind):
        rows, cols = np.shape(img_blk)
        # D = self.__dctmtx(rows)
        D = self.__dct_matrices[ind]
        dct_coeff = np.matmul(np.matmul(D, img_blk), np.transpose(D))
        return(dct_coeff)

    def entropyFilt(self, img):
        return(entropy(img, square(self.entropy_filt_kernel_sze)))

    def computeScore(self, weighted_local_entropy, T_max):
        # normalize weighted T max matrix
        min_val = weighted_local_entropy.min()
        weighted_T_Max = weighted_local_entropy - min_val
        max_val = weighted_local_entropy.max()
        weighted_T_Max = weighted_local_entropy / max_val

        score = np.median(weighted_local_entropy)
        return(score)

    def TransformedDomainRecursiveFilter_Horizontal(self, I, D, sigma):
        # Feedback Coefficient (Appendix of the paper)
        a = np.exp(-np.sqrt(2) / sigma)
        F = copy.deepcopy(I)
        V = a ** D
        rows, cols = np.shape(I)

        # Left --> Right Filter
        for i in range(1, cols):
            F[:, i] = F[:, i] + np.multiply(V[:, i], (F[:, i-1] - F[:, i]))

        # Right --> Left Filter
        for i in range(cols-2, 1, -1):
            F[:, i] = F[:, i] + np.multiply(V[:, i+1], (F[:, i + 1] - F[:, i]))

        return(F)

    def RF(self, img, joint_img):
        if(len(joint_img) == 0):
            joint_img = img
        joint_img = joint_img.astype('float64')
        joint_img = joint_img / 255

        if(len(np.shape(joint_img)) == 2):
            cols, rows = np.shape(joint_img)
            channels = 1
        elif(len(np.shape(joint_img)) == 3):
            cols, rows, channels = np.shape(joint_img)
        # Estimate horizontal and vertical partial derivatives using finite differences.
        dIcdx = np.diff(joint_img, n=1, axis=1)
        dIcdy = np.diff(joint_img, n=1, axis=0)

        dIdx = np.zeros((cols, rows))
        dIdy = np.zeros((cols, rows))

        # Compute the l1 - norm distance of neighbor pixels.
        dIdx[:, 1::] = abs(dIcdx)
        dIdy[1::, :] = abs(dIcdy)

        dHdx = (1 + self.sigma_s_RF_filter / self.sigma_r_RF_filter * dIdx)
        dVdy = (1 + self.sigma_s_RF_filter / self.sigma_r_RF_filter * dIdy)

        dVdy = np.transpose(dVdy)
        N = self.num_iterations_RF_filter
        F = copy.deepcopy(img)
        for i in range(self.num_iterations_RF_filter):
            # Compute the sigma value for this iteration (Equation 14 of our paper).
            sigma_H_i = self.sigma_s_RF_filter * \
                np.sqrt(3) * 2 ** (N - (i + 1)) / np.sqrt(4 ** N - 1)
            F = self.TransformedDomainRecursiveFilter_Horizontal(
                F, dHdx, sigma_H_i)
            F = np.transpose(F)

            F = self.TransformedDomainRecursiveFilter_Horizontal(
                F, dVdy, sigma_H_i)
            F = np.transpose(F)

        return(F)

    def detectBlur(
        self, 
        img 
    ):
        ori_rows, ori_cols = np.shape(img)
        # perform initial gausssian smoothing
        InputImageGaus = cv2.GaussianBlur(img, (3, 3), sigmaX=0.5, sigmaY=0.5)
        __gradient_image = self.computeImageGradientMagnitude(InputImageGaus)

        total_num_layers = 1 + sum(self.scales)

        # create all dct_matrices beforehand to save computation time
        self.__createDCT_Matrices()

        # Create Frequency Labels at all the scalesv
        self.__computeFrequencyBands()

        # Compute the indices of the high frequency content inside each frequency band
        for i in range(self.num_scales):
            curr_freq_band = self.__freqBands[i]
            self.freq_index.append(np.where(curr_freq_band == 0))

        __padded_image = np.pad(__gradient_image, int(
            np.floor(max(self.scales)/2)), mode='constant')

        rows, cols = np.shape(__padded_image)
        L = []

        total_num_points = len([i for i in range(int(max(self.scales)/2), rows - int(max(self.scales)/2), self.downsampling_factor)]) * \
            len([j for j in range(int(max(self.scales) / 2), cols -
                int(max(self.scales) / 2), self.downsampling_factor)])
        L = np.zeros((total_num_points, total_num_layers))

        iter = 0
        n = 0
        old_progress = 0
        for i in range(int(max(self.scales)/2), rows - int(max(self.scales)/2), self.downsampling_factor):
            #old_progress = self.disp_progress(i, rows, old_progress)
            m = 0
            n += 1
            for j in range(int(max(self.scales) / 2), cols - int(max(self.scales) / 2), self.downsampling_factor):
                m += 1
                high_freq_components = []
                for ind, curr_scale in enumerate(self.scales):
                    Patch = __padded_image[i-np.int(curr_scale/2): i+np.int(
                        curr_scale/2) + 1, j-np.int(curr_scale/2): j+np.int(curr_scale/2) + 1]
                    dct_coefficients = np.abs(
                        self.__getDCTCoefficients(Patch, ind))

                    # store all high frequency components
                    high_freq_components.append(
                        dct_coefficients[self.freq_index[ind]])

                # Find the first `total_num_layers` smallest values in all the high frequency components - we must not sort the entire array since that is very inefficient
                high_freq_components = np.hstack(high_freq_components)
                result = np.argpartition(
                    high_freq_components, total_num_layers)
                L[iter, :] = high_freq_components[result[:total_num_layers]]
                iter += 1

        L = np.array(L)

        # normalize the L matrix
        for i in range(total_num_layers):
            max_val = max(L[:, i])
            L[:, i] = L[:, i] / max_val

        # perform max pooling on the normalized frequencies
        ind1d = 0
        T_max = np.zeros((n, m))
        max_val = 0
        min_val = 99999
        for i in range(n):
            for j in range(m):
                T_max[i][j] = max(L[ind1d, :])
                max_val = max(max_val, T_max[i][j])
                min_val = min(min_val, T_max[i][j])
                ind1d += 1

        # Final Map and Post Processing
        local_entropy = self.entropyFilt(T_max)
        weighted_local_entropy = np.multiply(local_entropy, T_max)

        score = self.computeScore(weighted_local_entropy, T_max)
        rows, cols = np.shape(weighted_local_entropy)

        # resize the input image to match the size of local_entropy matrix
        resized_input_image = cv2.resize(InputImageGaus, (cols, rows))
        aSmooth = cv2.GaussianBlur(
            resized_input_image, (3, 3), sigmaX=1, sigmaY=1)
        final_map = self.RF(weighted_local_entropy, aSmooth)

        # resize the map to the original resolution
        final_map = cv2.resize(final_map, (ori_cols, ori_rows))

        # normalize the map
        final_map = final_map / np.max(final_map)
        return(score,final_map)
