import matplotlib.pyplot as plt
import cv2
import numpy as np
import time

cameraman = cv2.imread('cameraman.png')
fognoise = cv2.imread('fognoise.png')
fog = cv2.imread('fog.png')
frostfog = cv2.imread('frostfog.png')
bat = cv2.imread('bat.png')
tree = cv2.imread('tree.png', 0)
treeM = cv2.imread('treeM.png', 0)


def draw_histograms_cameraman():
    my_dpi = 240
    plt.figure(figsize=(1024 / my_dpi, 512 / my_dpi), dpi=my_dpi)

    hist, bin_edges = np.histogram(cameraman.ravel(), bins=256, range=[0, 256], normed=True)
    F1 = np.cumsum(hist)
    width = 1 / 1.5
    plt.bar(np.arange(0, 256), hist, width, color='red', alpha=0.5)
    plt.bar(np.arange(0, 256) + width, F1 * max(hist), width, color='blue', alpha=0.5)
    plt.title('Cameraman Histogram')
    plt.savefig('Cameraman_histograms.png', dpi=my_dpi)
    plt.show()


def draw_histograms_bat():
    my_dpi = 240
    plt.figure(figsize=(1024 / my_dpi, 512 / my_dpi), dpi=my_dpi)
    histbat, bin_edges = np.histogram(bat.ravel(), bins=256, range=[0, 256], normed=True)

    F1 = np.cumsum(histbat)
    width = 1 / 1.5
    plt.bar(np.arange(0, 256), histbat, width, color='red', alpha=0.5)
    plt.bar(np.arange(0, 256) + width, F1 * max(histbat), width, color='blue', alpha=0.5)
    plt.title('Bat Histogram')
    plt.savefig('Bat_histograms.png', dpi=my_dpi)

    plt.show()


def draw_histograms_fog():
    my_dpi = 240
    plt.figure(figsize=(1024 / my_dpi, 512 / my_dpi), dpi=my_dpi)
    hist, bin_edges = np.histogram(fog.ravel(), bins=256, range=[0, 256], normed=True)
    F1 = np.cumsum(hist)
    width = 1 / 1.5
    plt.bar(np.arange(0, 256), hist, width, color='red', alpha=0.5)
    plt.bar(np.arange(0, 256) + width, F1 * max(hist), width, color='blue', alpha=0.5)
    plt.title('Fog Histogram')
    plt.savefig('Fog_histograms.png', dpi=my_dpi)

    plt.show()


def draw_histograms_fogNoise():
    my_dpi = 240
    plt.figure(figsize=(1024 / my_dpi, 512 / my_dpi), dpi=my_dpi)
    bins = len(fognoise)
    print(bins)
    hist, bin_edges = np.histogram(fognoise.ravel(), bins=256, range=[0, 256], normed=True)
    F1 = np.cumsum(hist)
    width = 1 / 1.5
    plt.bar(np.arange(0, 256), hist, width, color='red', alpha=0.5)
    plt.bar(np.arange(0, 256) + width, F1 * max(hist), width, color='blue', alpha=0.5)
    plt.title('FogNoise Histogram')
    plt.savefig('FogNoise_histograms.png', dpi=my_dpi)
    plt.show()


# Mean Verses Gaussian

def mean_vs_gaussian():
    # preparing kernel
    kernel = np.ones((5, 5), np.float32) / 25
    # applying 5x5 mean filter
    dst = cv2.filter2D(cameraman, -1, kernel)

    # applying 5x5 gaussian filter
    blur = cv2.GaussianBlur(cameraman, (5, 5), 0)
    # plotting images after applying both filters
    plt.subplot(221), plt.imshow(blur), plt.title('Gaussian')
    plt.xticks([]), plt.yticks([])
    plt.subplot(222), plt.imshow(dst), plt.title('Mean Filter')
    plt.xticks([]), plt.yticks([])
    # Plotting histograms
    plt.subplot(223), plt.hist(blur.flatten(), 256, range=(0, 255), fc='w', ec='k', normed=True), plt.title(
        'Gaussian Histogram')
    plt.subplot(224), plt.hist(dst.flatten(), 256, range=(0, 255), fc='w', ec='k', normed=True), plt.title(
        'Mean Histogram')
    # for the image to be full screen
    figManager = plt.get_current_fig_manager()
    figManager.window.showMaximized()
    plt.show()


def selective_Median_Filter():
    # measuring runtime for median filter on all pixels
    start_time = time.clock()
    median = median_filter()
    print('Pre-enhancement runtime : ', "--- %s seconds ---" % (time.clock() - start_time))
    plt.imshow(median)
    plt.show()
    # measuring runtime for median filter on too black and woo white pixels only(noise)
    start_time2 = time.clock()
    median2 = selective_median_filter()
    print('Post-enhancement runtime : ', "--- %s seconds ---" % (time.clock() - start_time2))
    plt.imshow(median2)
    plt.show()


def median_filter():
    img = cv2.cvtColor(fognoise, cv2.COLOR_BGR2GRAY)
    # neglecting corners and edges of image
    for i in range(2, img.shape[0] - 2):
        for j in range(2, img.shape[1] - 2):
            # apply filter on all pixels
            sorted = np.sort(img[i - 2:i + 3, j - 2:j + 3].flatten())
            idx = (len(sorted) / 2) + 1 if len(sorted) % 2 == 0 else (len(sorted)) / 2
            img[i][j] = sorted[int(idx)]
    img2 = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    return img2


def selective_median_filter():
    img = cv2.cvtColor(fognoise, cv2.COLOR_BGR2GRAY)
    # neglecting corners and edges of image
    for i in range(2, img.shape[0] - 2):
        for j in range(2, img.shape[1] - 2):
            p = img[i][j]  # apply it on too black and too white only
            if 0 <= p <= 30 or 225 <= p <= 255:
                sorted = np.sort(img[i - 2:i + 3, j - 2:j + 3].flatten())
                idx = (len(sorted) / 2) + 1 if len(sorted) % 2 == 0 else (len(sorted)) / 2
                img[i][j] = sorted[int(idx)]
    img2 = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    return img2


def contrastStretching_histogramEqualization():
    # contrast stretching
    # normalize the image with type NORM_MINMAX
    img_rescale = frostfog
    cv2.normalize(img_rescale, dst=img_rescale, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX)

    # histogram equalization
    # convert the image back to grey scale format
    img_gray = cv2.cvtColor(frostfog, cv2.COLOR_BGR2GRAY)
    img_gray = cv2.equalizeHist(img_gray)
    # convert the image back to RGB format
    img_output = cv2.cvtColor(img_gray, cv2.COLOR_GRAY2BGR)

    # drawing images after contrast stretching and histogram equalization
    plt.subplot(221), plt.imshow(img_rescale), plt.title('Contrast stretching')
    plt.xticks([]), plt.yticks([])
    plt.subplot(222), plt.imshow(img_output), plt.title('Histogram equalization')
    plt.xticks([]), plt.yticks([])

    # calculating and plotting histograms for both
    plt.subplot(223), plt.hist(img_rescale.flatten(), 256, range=[0, 256], fc='r', ec='k', normed=True), plt.title(
        'Contrast stretching Histogram')
    # plt.xticks([]), plt.yticks([])

    # check the cumulative if it is linear
    plt.subplot(224), plt.hist(img_output.flatten(), 256, range=[0, 256], fc='b', ec='k', normed=True), plt.title(
        'Histogram equalization Histogram')
    # plt.xticks([]), plt.yticks([])

    # for the image to be full screen
    figManager = plt.get_current_fig_manager()
    figManager.window.showMaximized()
    plt.show()