import math
import random

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

counter = 1


def turnToBinary(img):
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    Binary = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
    return Binary


def segmented_linear_transformation(img):
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    hist = cv2.calcHist([img], [0], None, [256], [0, 256])
    left = -1
    right = 256
    for i in range(1, 256):
        hist[i] = hist[i - 1] + hist[i]
    for i in range(0, 256):
        if hist[i] >= hist[255] * 0.1 and left == -1:
            left = i
        elif hist[i] >= hist[255] * 0.9 and right == 256:
            right = i
    h, w = img.shape
    out = np.zeros(img.shape, np.uint8)
    for i in range(h):
        for j in range(w):
            pix = img[i][j]
            if pix < left:
                out[i][j] = 0.5 * pix
            elif left <= pix < right:
                out[i][j] = 3.6 * pix - 310
            else:
                out[i][j] = 0.238 * pix + 194
    out = np.around(out)
    out = out.astype(np.uint8)
    return out


def hsv_ajust(img, k, t):
    img = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            tmp = float(img[i][j][t] * k)
            if tmp > 255:
                tmp = 255
            img[i][j][t] = int(tmp)
    img = cv2.cvtColor(img, cv2.COLOR_HSV2RGB)
    return img


def size_change(img, x, y, z):  # x weight y height z interpolation
    itp = {'最近邻插值': cv2.INTER_NEAREST, '双线性插值': cv2.INTER_LINEAR, '基于局部像素的重采样': cv2.INTER_AREA,
           '基于4x4像素邻域的3次插值法': cv2.INTER_CUBIC, '基于8x8像素邻域的Lanczos插值': cv2.INTER_LANCZOS4}
    return cv2.resize(img, (int(y), int(x)), itp.get(z))


def image_panning(img, x, y):  # 图片平移
    M = np.float32([[1, 0, y], [0, 1, x]])
    return cv2.warpAffine(img, M, (img.shape[1], img.shape[0]))


def image_rotate(img, x):
    height, width, channel = img.shape
    M = cv2.getRotationMatrix2D((width / 2, height / 2), x, 1)
    return cv2.warpAffine(img, M, (width, height))


def get_gray_hist(img):
    global counter
    filename = 'src/' + str(counter) + '_gray.jpg'
    counter = counter + 1
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    plt.figure(counter)
    plt.hist(img.ravel(), 256)
    plt.savefig(filename)
    return cv2.cvtColor(cv2.imread(filename, 1), cv2.COLOR_BGR2RGB)


def get_colorful_hist(img):
    global counter
    filename = 'src/' + str(counter) + '_colorful.jpg'
    counter = counter + 1
    plt.figure(counter)
    color = ['b', 'g', 'r']
    for index, c in enumerate(color):
        hist = cv2.calcHist([img], [index], None, [256], [0, 256])
        print(hist)
        plt.plot(hist, color=c)
        plt.xlim([0, 256])
    plt.savefig(filename)
    return cv2.cvtColor(cv2.imread(filename, 1), cv2.COLOR_BGR2RGB)


def Hough_transform(img):
    img = cv2.GaussianBlur(img, (3, 3), 0)  # 高斯模糊
    edges = cv2.Canny(img, 50, 150, apertureSize=3)  # Canny边缘检测
    lines = cv2.HoughLinesP(edges, 1, np.pi / 180, 30)
    img_empty = np.zeros((img.shape[0], img.shape[1], 3), np.uint8)
    for i in lines:
        for x1, y1, x2, y2 in i:
            cv2.line(img_empty, (x1, y1), (x2, y2), (255, 255, 255), 1)
    return img_empty


def gradient_img_sharpen(img):
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    h, w = img.shape
    img.astype('float')
    gradient = np.zeros([h, w], dtype='float')
    for x in range(h):
        for y in range(w):
            if x == h - 1:
                gx = abs(float(img[x][y]) - float(img[x - 1][y]))
            else:
                gx = abs(float(img[x + 1][y]) - float(img[x][y]))
            if y == w - 1:
                gy = abs(float(img[x][y]) - float(img[x][y - 1]))
            else:
                gy = abs(float(img[x][y + 1]) - float(img[x][y]))
            gradient[x][y] = float(gx) + float(gy)
    sharp = gradient
    sharp = np.where(sharp > 255, 255, sharp)
    sharp = np.where(sharp < 0, 0, sharp)
    sharp = sharp.astype('uint8')
    return sharp


def Roberts_img_sharpen(img):
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    kernelx = np.array([[-1, 0], [0, 1]], dtype=int)
    kernely = np.array([[0, -1], [1, 0]], dtype=int)
    x = cv2.filter2D(img, cv2.CV_16S, kernelx)
    y = cv2.filter2D(img, cv2.CV_16S, kernely)
    absX = cv2.convertScaleAbs(x)
    absY = cv2.convertScaleAbs(y)
    Roberts = cv2.addWeighted(absX, 0.5, absY, 0.5, 0)
    return Roberts


def Prewitt_img_sharpen(img):
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    kernelx = np.array([[1, 0, -1], [1, 0, -1], [1, 0, -1]], dtype=int)
    kernely = np.array([[-1, -1, -1], [0, 0, 0], [1, 1, 1]], dtype=int)
    x = cv2.filter2D(img, cv2.CV_16S, kernelx)
    y = cv2.filter2D(img, cv2.CV_16S, kernely)
    absX = cv2.convertScaleAbs(x)
    absY = cv2.convertScaleAbs(y)
    Prewitt = cv2.addWeighted(absX, 0.5, absY, 0.5, 0)
    return Prewitt


def Sobel_img_sharpen(img):
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    x = cv2.Sobel(img, cv2.CV_16S, 1, 0)
    y = cv2.Sobel(img, cv2.CV_16S, 0, 1)
    absX = cv2.convertScaleAbs(x)
    absY = cv2.convertScaleAbs(y)
    Sobel = cv2.addWeighted(absX, 0.5, absY, 0.5, 0)
    return Sobel


def Laplacian_img_sharpen(img):
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    gray = cv2.GaussianBlur(img, (5, 5), 0)
    dst = cv2.Laplacian(gray, cv2.CV_16S, ksize=3)
    Laplacian = cv2.convertScaleAbs(dst)
    return Laplacian


def LoG_img_sharpen(img):
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    gray = cv2.copyMakeBorder(img, 2, 2, 2, 2, borderType=cv2.BORDER_REPLICATE)
    gray = cv2.GaussianBlur(gray, (3, 3), 0, 0)
    LoG = np.array([[0, 0, -1, 0, 0], [0, -1, -2, -1, 0], [-1, -2, 16, -2, -1], [0, -1, -2, -1, 0], [0, 0, -1, 0, 0]])
    h = gray.shape[0]
    w = gray.shape[1]
    image = np.zeros([h, w], dtype=int)
    for i in range(2, h - 2):
        for j in range(2, w - 2):
            image[i, j] = np.sum(LoG * gray[i - 2:i + 3, j - 2:j + 3, 1])
    image = cv2.convertScaleAbs(image)
    return image


def Canny_img_sharpen(img):
    img = cv2.GaussianBlur(img, (3, 3), 0)
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gx = cv2.Sobel(img, cv2.CV_16SC1, 1, 0)
    gy = cv2.Sobel(img, cv2.CV_16SC1, 0, 1)
    edge_output = cv2.Canny(gx, gy, 50, 100)
    return edge_output


def neighberhood_average_smooth(img):
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    h, w = img.shape
    out = np.zeros((h, w))
    for i in range(h):
        for j in range(w):
            if i == 0 or j == 0 or i == h - 1 or j == w - 1:
                out[i][j] = img[i][j]
            else:
                ans = 0
                for k in range(-1, 2):
                    for s in range(-1, 2):
                        ans = ans + img[i + k][j + s]
                ans = ans / 9
                out[i][j] = ans
    return out


def median_filtering_smooth(img):
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    h, w = img.shape
    out = np.zeros((h, w))
    for i in range(h):
        for j in range(w):
            if i == 0 or j == 0 or i == h - 1 or j == w - 1:
                out[i][j] = img[i][j]
            else:
                pix = []
                for k in range(-1, 2):
                    for s in range(-1, 2):
                        pix.append(img[i + k][j + s])
                pix.sort()
                out[i][j] = pix[4]
    return out


def Ideal_low_pass_filtering_smooth(img):
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    img = np.float32(img)
    dft = cv2.dft(img, flags=cv2.DFT_COMPLEX_OUTPUT)
    dft = np.fft.fftshift(dft)
    h, w = img.shape
    mask = np.zeros((h, w, 2), dtype=np.uint8)
    mask[int(h / 2) - 20:int(h / 2) + 20, int(w / 2) - 20:int(w / 2) + 20] = 1
    f = dft * mask
    f = np.fft.ifftshift(f)
    img = cv2.idft(f)
    img = cv2.magnitude(img[:, :, 0], img[:, :, 1])
    global counter
    filename = 'src/' + str(counter) + '_ideal_low_pass.png'
    plt.imsave(filename, img, cmap='gray')
    counter = counter + 1
    img = cv2.imread(filename, 1)
    return img


def ButterWorth_filtering_smooth(img):
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    img = np.float32(img)
    dft = cv2.dft(img, flags=cv2.DFT_COMPLEX_OUTPUT)
    dft = np.fft.fftshift(dft)
    h, w = img.shape
    mask = np.zeros((h, w, 2), dtype=np.float32)
    for i in range(h):
        for j in range(w):
            x = i - h / 2
            y = j - w / 2
            D = np.sqrt(x ** 2 + y ** 2)
            k = D / 20
            k = k ** 4
            k = k + 1.0
            k = 1.0 / k
            mask[i, j, 0] = mask[i, j, 1] = k
    f = dft * mask
    f = np.fft.ifftshift(f)
    img = cv2.idft(f)
    img = cv2.magnitude(img[:, :, 0], img[:, :, 1])
    global counter
    filename = 'src/' + str(counter) + '_butterworth_low_pass.png'
    plt.imsave(filename, img, cmap='gray')
    counter = counter + 1
    img = cv2.imread(filename, 1)
    return img


def Gaussian_filtering_smooth(img):
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    img = np.float32(img)
    dft = cv2.dft(img, flags=cv2.DFT_COMPLEX_OUTPUT)
    dft = np.fft.fftshift(dft)
    h, w = img.shape
    mask = np.zeros((h, w, 2), dtype=np.float32)
    for i in range(h):
        for j in range(w):
            x = i - h / 2
            y = j - w / 2
            D = np.sqrt(x ** 2 + y ** 2)
            k = D / 20
            k = k * k
            k = k * - 0.5
            k = 2.7182818 ** k
            mask[i][j][0] = mask[i][j][1] = k
    f = dft * mask
    f = np.fft.ifftshift(f)
    img = cv2.idft(f)
    img = cv2.magnitude(img[:, :, 0], img[:, :, 1])
    global counter
    filename = 'src/' + str(counter) + '_gaussian_low_pass.png'
    plt.imsave(filename, img, cmap='gray')
    counter = counter + 1
    img = cv2.imread(filename, 1)
    return img


def produce_pepper_salt_noise(img):
    h, w, t = img.shape
    for i in range(h):
        for j in range(w):
            rnd = random.random()
            if rnd < 0.05:
                img[i][j] = [0, 0, 0]
            elif rnd > 0.95:
                img[i][j] = [255, 255, 255]
    return img


def produce_Gaussian_noise(img):
    img = np.array(img / 255, dtype=float)
    noise = np.random.normal(0, 0.1, img.shape)
    out = img + noise
    out = np.clip(out, 0.0, 1.0)
    out = np.uint8(out * 255)
    return out


def image_erode(img):
    if len(img.shape) == 3:
        img = turnToBinary(img)
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (5, 5), (-1, -1))
    img = cv2.erode(img, kernel)
    return img


def image_dilate(img):
    if len(img.shape) == 3:
        img = turnToBinary(img)
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (5, 5), (-1, -1))
    img = cv2.dilate(img, kernel)
    return img