# This Python file uses the following encoding: utf-8
import functools
import os
import sys
from math import sqrt
from pathlib import Path
from typing import Callable

import cv2
import matplotlib.pyplot as plt
import numpy as np
from PySide6.QtCore import Slot, QUrl, QPoint, QFileInfo
from PySide6.QtGui import QImage, QPainter, QColor
from PySide6.QtQml import QQmlApplicationEngine, QmlElement
from PySide6.QtQuick import QQuickPaintedItem
from PySide6.QtWidgets import QApplication
from numba import jit
from lapstyle import LapStylePredictor

QML_IMPORT_NAME = "Util"
QML_IMPORT_MAJOR_VERSION = 1
QML_IMPORT_MINOR_VERSION = 0


@QmlElement
class ImageItem(QQuickPaintedItem):
    matrices: list[np.ndarray] = []
    matrix: np.ndarray = None
    image: QImage = None
    xpos = 0
    ypos = 0

    def __init__(self, parent=None):
        super(ImageItem, self).__init__(parent)

    def paint(self, painter: QPainter):
        if self.image is None:
            painter.fillRect(self.boundingRect(), QColor("transparent"))
        else:
            x = self.width() / 2 - self.image.width() / 2
            if x < 0:
                x = 0
            y = self.height() / 2 - self.image.height() / 2
            if y < 0:
                y = 0
            painter.drawImage(x, y, self.image, self.xpos, self.ypos)

    def setImage(self, matrix: np.ndarray):
        self.matrix = matrix
        # 彩色图片
        if len(matrix.shape) > 2:
            height, width, channel = matrix.shape
            self.image = QImage(matrix.data, width, height, width * channel, QImage.Format_RGB888)
        # 灰度图片
        else:
            height, width = matrix.shape
            self.image = QImage(matrix.data, width, height, width, QImage.Format_Grayscale8)
        self.update()

    @Slot(float, float, result="QPoint")
    def getPoint(self, x, y):
        ox = np.round(x + self.xpos)
        oy = np.round(y + self.ypos)
        return QPoint(ox, oy)

    def appendMatrix(self):
        if self.matrix is None:
            return
        if len(self.matrices) >= 1 and np.array_equal(self.matrices[-1], self.matrix):
            return
        self.matrices.append(self.matrix)

    @Slot()
    def withdraw(self):
        if self.matrix is None:
            return
        if np.array_equal(self.matrix, self.matrices[-1]) and len(self.matrices) >= 2:
            self.matrices.pop()
        self.setImage(self.matrices[-1])

    @staticmethod
    def process(isGray=False):
        def decorator(func: Callable):
            @functools.wraps(func)
            def wrapper(self, *args):
                image = self.getMatrix(isGray)
                if image is None:
                    return
                out = func(self, *args, src=image)
                if out is not None:
                    self.setImage(out)
                    self.appendMatrix()

            return wrapper

        return decorator

    def getMatrix(self, isGray=False):
        if len(self.matrices) < 1:
            return None
        matrix = self.matrices[-1]
        if not isGray or len(matrix.shape) < 3:
            return matrix.copy()
        return cv2.cvtColor(matrix, cv2.COLOR_RGB2GRAY)

    @Slot(QUrl)
    def readFile(self, filepath: QUrl):
        image = cv2.imdecode(np.fromfile(filepath.toLocalFile(), dtype=np.uint8), cv2.IMREAD_COLOR)
        matrix = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        self.setImage(matrix)
        self.matrices.clear()
        self.appendMatrix()

    @Slot(QUrl)
    def saveFile(self, filepath: QUrl):
        if self.matrix is None:
            return
        if len(self.matrix.shape) > 2:
            img = cv2.cvtColor(self.matrix, cv2.COLOR_RGB2BGR)
        else:
            img = self.matrix

        fileInfo = QFileInfo(filepath.toLocalFile())
        cv2.imencode("." + fileInfo.suffix(), img)[1].tofile(filepath.toLocalFile())

    @Slot(int, int)
    def moveBy(self, x: int, y: int):
        if self.image is None:
            return
        if x + self.xpos + self.width() <= self.image.width():
            ox = x + self.xpos
        else:
            ox = self.image.width() - self.width()
        if ox < 0:
            ox = 0
        if y + self.ypos + self.height() <= self.image.height():
            oy = y + self.ypos
        else:
            oy = self.image.height() - self.height()
        if oy < 0:
            oy = 0
        if ox is not self.xpos or oy is not self.ypos:
            self.xpos = ox
            self.ypos = oy
            self.update()

    # 图像变换
    @Slot(int)
    @process()
    def removeChannel(self, channel, src):
        img = src
        img[..., channel] = 0
        return img

    @Slot(int, int, QColor)
    @process()
    def translation(self, dx, dy, border: QColor, src):
        img = src  # 读取彩色图像(BGR)
        rows, cols = img.shape[:2]
        mat = np.float32([[1, 0, dx], [0, 1, dy]])  # 构造平移变换矩阵
        (r, g, b, a) = border.getRgb()
        dst = cv2.warpAffine(img, mat, (cols, rows), borderValue=(r, g, b, a))
        return dst

    @Slot(float)
    @process()
    def centralRotation(self, theta, src):
        # 1.36 图像旋转 (以任意点 (x0,y0) 为中心旋转)
        img = src
        height, width = img.shape[:2]  # 图片的高度和宽度
        x0, y0 = width // 2, height // 2  # 以图像中心作为旋转中心
        mat = cv2.getRotationMatrix2D((x0, y0), theta, 1.0)
        dst = cv2.warpAffine(img, mat, (width, height))  # 设置白色填充
        return dst

    @Slot(int)
    @process()
    def rightAngleRotation(self, angle, src):
        img = src
        img_r = None
        if angle == 1:
            img_r = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
        elif angle == 2:
            img_r = cv2.rotate(img, cv2.ROTATE_180)
        if angle == 3:
            img_r = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
        return img_r

    @Slot(int)
    @process()
    def flip(self, method, src):
        img = src
        img_flip = None
        if method == 1:
            img_flip = cv2.flip(img, 0)  # 垂直翻转
        if method == 2:
            img_flip = cv2.flip(img, 1)  # 水平翻转
        if method == 3:
            img_flip = cv2.flip(img, -1)  # 水平和垂直翻转
        return img_flip

    @Slot(float, float)
    @process()
    def resize(self, width_size, height_size, src):
        img = src
        dst = cv2.resize(img, None, fx=width_size, fy=height_size, interpolation=cv2.INTER_AREA)
        return dst

    @Slot()
    @process()
    def affine(self, src):
        image = src
        rows, cols = image.shape[:2]
        pts1 = np.float32([[50, 50], [200, 50], [50, 200]])  # 初始位置
        pts2 = np.float32([[50, 100], [200, 50], [100, 250]])  # 终止位置
        mtr = cv2.getAffineTransform(pts1, pts2)  # 计算 2x3 变换矩阵 MA
        dst = cv2.warpAffine(image, mtr, (cols, rows))  # 实现仿射变换
        return dst

    @Slot()
    @process()
    def shear(self, src):
        # 1.41 图像的错切
        img = src
        height, width = img.shape[:2]  # 图片的高度和宽度
        mat = np.float32([[1, 0.2, 0], [0, 1, 0]])  # 构造错切变换矩阵
        img_shear = cv2.warpAffine(img, mat, (width, height))
        return img_shear

    @Slot()
    @process()
    def projective(self, src):
        img = src
        h, w = img.shape[:2]  # 图片的高度和宽度

        point_src = np.float32([[0, 0], [w, 0], [0, h], [w, h]])  # 原始图像中 4点坐标
        point_dst = np.float32(
            [[int(w / 3), int(h / 3)], [int(w * 2 / 3), int(h / 3)], [0, h], [w, h]])  # 变换图像中 4点坐标

        mp = cv2.getPerspectiveTransform(point_src, point_dst)  # 计算投影变换矩阵 M
        img_p = cv2.warpPerspective(img, mp, (w, h), flags=cv2.INTER_AREA, borderMode=cv2.BORDER_WRAP)
        return img_p

    # 图形绘制
    @Slot(QPoint, QPoint, QColor, int)
    @process()
    def line(self, start: QPoint, end: QPoint, color: QColor, width, src):
        img = src
        if start.isNull() or end.isNull():
            return
        (r, g, b, a) = color.getRgb()
        cv2.line(img, start.toTuple(), end.toTuple(), (r, g, b, a), width)
        return img

    @Slot(QPoint, QPoint, QColor, int)
    @process()
    def rectangle(self, start: QPoint, end: QPoint, color: QColor, width, src):
        img = src
        if start.isNull() or end.isNull():
            return
        (r, g, b, a) = color.getRgb()
        cv2.rectangle(img, start.toTuple(), end.toTuple(), (r, g, b, a), width)
        return img

    @Slot(QPoint, QPoint, QColor, int)
    @process()
    def circle(self, center: QPoint, side: QPoint, color: QColor, width, src):
        img = src
        if center.isNull() or side.isNull():
            return
        (center_x, center_y) = center.toTuple()
        (side_x, side_y) = side.toTuple()
        radius = np.round(sqrt(np.power(center_x - side_x, 2) + np.power(center_y - side_y, 2)))
        (r, g, b, a) = color.getRgb()
        cv2.circle(img, center.toTuple(), int(radius), (r, g, b, a), width)
        return img

    @Slot(QPoint, QColor, int, int, str)
    @process()
    def text(self, start: QPoint, color: QColor, fontsize, width, text, src):
        img = src
        if start.isNull():
            return
        (r, g, b, a) = color.getRgb()
        cv2.putText(img, text, start.toTuple(), cv2.FONT_HERSHEY_SCRIPT_COMPLEX, fontsize, (r, g, b, a), width,
                    cv2.LINE_AA, 0)
        return img

    # 灰度运算

    @Slot(int)
    @process(isGray=True)
    def binaryTransformation(self, thresh, src):
        img_gray = src
        if img_gray is None:
            return
        _, img = cv2.threshold(img_gray, thresh, 255, cv2.THRESH_BINARY)  # 转换为二值图像, thresh=191
        return img

    @Slot()
    @process(isGray=True)
    def linearGrayscaleTransformation(self, src):
        # 通过遍历对不同像素范围内进行分段线性变化，在这里分三段函数进行分段线性变化, 主要还是考察for循环的应用
        # y=0.5*x(x<50)
        # y=3.6*x-310(50<=x<150)
        # y=0.238*x+194(x>=150)
        img = src

        @jit(nopython=True)
        def transform(image):
            out = np.zeros(image.shape, np.uint8)
            for pos, pix in np.ndenumerate(image):
                if pix < 50:
                    out[pos] = 0.5 * pix
                elif pix < 150:
                    out[pos] = 3.6 * pix - 310
                else:
                    out[pos] = 0.238 * pix + 194
            return out

        res = transform(img)
        return res

    @Slot(int, int)
    @process(isGray=True)
    def grayLayered(self, low_thresh, high_thresh, src):
        # # 1.53 分段线性灰度变换 (灰度级分层)  # Gray layered
        img_gray = src
        '''
        另一种方法
        img_layer1 = img_gray.copy()
        img_layer1[(img_layer1[:, :] < low_thresh) | (img_layer1[:, :] > high_thresh)] = 0  # 其它区域：黑色
        img_layer1[(img_layer1[:, :] >= low_thresh) & (img_layer1[:, :] <= high_thresh)] = 255  # 灰度级窗口：白色
        '''

        img_layer = img_gray.copy()
        img_layer[(img_layer[:, :] >= low_thresh) & (img_layer[:, :] <= high_thresh)] = 255  # 灰度级窗口：白色，其它区域不变

        return img_layer

    @Slot()
    @process(isGray=True)
    def logTrans(self, src):
        # 1.55 图像的非线性灰度变换：对数变换
        img = src
        img_log = np.log(img + 1.0)
        cv2.normalize(img_log, img_log, 0, 255, cv2.NORM_MINMAX)
        img_log_abs = cv2.convertScaleAbs(img_log)
        return img_log_abs

    @Slot(float)
    @process(isGray=True)
    def gammaTrans(self, gamma, src):
        # 1.56 图像的非线性灰度变换: 幂律变换 (伽马变换)
        img = src
        img_gamma = np.power(img, gamma)
        cv2.normalize(img_gamma, img_gamma, 0, 255, cv2.NORM_MINMAX)
        img_gamma_abs = cv2.convertScaleAbs(img_gamma)
        return img_gamma_abs

    @Slot()
    @process(isGray=True)
    def grayHistogram(self, src):
        # 1.57 图像的灰度直方图
        img = src
        hist_cv = cv2.calcHist([img], [0], None, [256], [0, 256])  # OpenCV 函数 cv2.calcHist
        plt.bar(range(256), hist_cv[:, 0])
        '''
        用opencv和numpy两种方法得到直方图
        hist_np, bins = np.histogram(img.flatten(), 256)
        plt.subplot(132,xticks=[], yticks=[])
        plt.axis([0,255,0,np.max(hist_cv)])
        plt.bar(range(256), hist_cv[:,0])
        plt.title("Gray Hist(cv2.calcHist)")
        plt.subplot(133,xticks=[], yticks=[])
        plt.axis([0,255,0,np.max(hist_cv)])
        plt.bar(bins[:-1], hist_np)
        plt.title("Gray Hist(np.histogram)")
        '''
        plt.show()

    # 边缘检测
    @Slot()
    @process(isGray=True)
    def roberts(self, src):
        img = src
        # Roberts 边缘算子
        kernel_Roberts_x = np.array([[1, 0], [0, -1]])
        kernel_Roberts_y = np.array([[0, -1], [1, 0]])
        imgRoberts_x = cv2.filter2D(img, cv2.CV_16S, kernel_Roberts_x)
        imgRoberts_y = cv2.filter2D(img, cv2.CV_16S, kernel_Roberts_y)
        imgRoberts_x_abs = cv2.convertScaleAbs(imgRoberts_x)
        imgRoberts_y_abs = cv2.convertScaleAbs(imgRoberts_y)
        imgRoberts = cv2.addWeighted(imgRoberts_x_abs, 0.5, imgRoberts_y_abs, 0.5, 0)
        return imgRoberts

    @Slot()
    @process(isGray=True)
    def sobel(self, src):
        img = src
        # kernel_Sobel_x = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
        # kernel_Sobel_y = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])
        imgSobel_x = cv2.Sobel(img, cv2.CV_16S, 1, 0)
        imgSobel_y = cv2.Sobel(img, cv2.CV_16S, 0, 1)
        imgSobel_x_abs = cv2.convertScaleAbs(imgSobel_x)
        imgSobel_y_abs = cv2.convertScaleAbs(imgSobel_y)
        imgSobel = cv2.addWeighted(imgSobel_x_abs, 0.5, imgSobel_y_abs, 0.5, 0)
        return imgSobel

    @Slot()
    @process(isGray=True)
    def prewitt(self, src):
        img = src
        kernel_Prewitt_x = np.array([[-1, 0, 1], [-1, 0, 1], [-1, 0, 1]])
        kernel_Prewitt_y = np.array([[1, 1, 1], [0, 0, 0], [-1, -1, -1]])
        imgPrewitt_x = cv2.filter2D(img, cv2.CV_16S, kernel_Prewitt_x)
        imgPrewitt_y = cv2.filter2D(img, cv2.CV_16S, kernel_Prewitt_y)
        imgPrewitt_x_abs = cv2.convertScaleAbs(imgPrewitt_x)
        imgPrewitt_y_abs = cv2.convertScaleAbs(imgPrewitt_y)
        imgPrewitt = cv2.addWeighted(imgPrewitt_x_abs, 0.5, imgPrewitt_y_abs, 0.5, 0)
        return imgPrewitt

    @Slot()
    @process(isGray=True)
    def laplacian(self, src):
        img = src
        # kernel_Laplacian_K1 = np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]])
        # imgLaplacian = cv2.filter2D(img, -1, kernel_Laplacian_K1)
        imgLaplacian = cv2.Laplacian(img, cv2.CV_16S, ksize=3)
        imgLaplacian_abs = cv2.convertScaleAbs(imgLaplacian)
        return imgLaplacian_abs

    @Slot(int, float)
    @process(isGray=True)
    def loG(self, size, sigma, src):
        img = src
        imgGaussBlur = cv2.GaussianBlur(img, (size, size), sigma)
        imgLoG = cv2.Laplacian(imgGaussBlur, cv2.CV_16S, ksize=size)
        imgLoG_abs = cv2.convertScaleAbs(imgLoG)
        return imgLoG_abs

    @Slot(int, int, int)
    @process(isGray=True)
    def canny(self, size, low_thresh, high_thresh, src):
        img = src

        kSize = (size, size)
        imgGauss = cv2.GaussianBlur(img, kSize, sigmaX=1.0)  # sigma=1.0
        sobel_x = cv2.Sobel(imgGauss, cv2.CV_16S, 1, 0)
        sobel_y = cv2.Sobel(imgGauss, cv2.CV_16S, 0, 1)

        imgCanny = cv2.Canny(sobel_x, sobel_y, low_thresh, high_thresh)
        return imgCanny

    # 图像平滑
    @Slot(int)
    @process()
    def blur(self, size, src):
        img = src
        imgBlur = cv2.blur(img, (size, size))
        return imgBlur

    @Slot(int)
    @process()
    def medianBlur(self, size, src):
        # 1.73：图像的非线性滤波 (中值滤波器)
        img = src
        imgMedianBlur = cv2.medianBlur(img, size)
        return imgMedianBlur

    @Slot(int)
    @process(isGray=True)
    def lowPassIdealBlur(self, D0, src):
        img = src
        height, width = img.shape[:2]
        centerX, centerY = int(height / 2), int(width / 2)
        x, y = np.mgrid[0:height, 0:width]
        D = np.sqrt(np.power(x - centerX, 2) + np.power(y - centerY, 2))
        H = np.where(D > D0, 0, 1)
        fft = np.fft.fft2(img)  # 傅里叶变换
        fft_shift = np.fft.fftshift(fft)  # 中心化
        fft_mask = fft_shift * H
        ifft_shift = np.fft.ifftshift(fft_mask)
        ifft = np.fft.ifft2(ifft_shift)
        ifft_abs = np.abs(ifft)
        return cv2.convertScaleAbs(ifft_abs)

    @Slot(int, int)
    @process(isGray=True)
    def lowPassButterworthBlur(self, D0, n, src):
        img = src
        height, width = img.shape[:2]
        centerX, centerY = int(height / 2), int(width / 2)
        x, y = np.mgrid[0:height, 0:width]
        D = np.sqrt(np.power(x - centerX, 2) + np.power(y - centerY, 2))
        H = 1.0 / (1.0 + np.power(D / D0, 2 * n))
        fft = np.fft.fft2(img)  # 傅里叶变换
        fft_shift = np.fft.fftshift(fft)  # 中心化
        fft_mask = fft_shift * H
        ifft_shift = np.fft.ifftshift(fft_mask)
        ifft = np.fft.ifft2(ifft_shift)
        ifft_abs = np.abs(ifft)
        return cv2.convertScaleAbs(ifft_abs)

    @Slot(int)
    @process(isGray=True)
    def lowPassGaussianBlur(self, D0, src):
        img = src
        height, width = img.shape[:2]
        centerX, centerY = int(height / 2), int(width / 2)
        x, y = np.mgrid[0:height, 0:width]
        D = np.sqrt(np.power(x - centerX, 2) + np.power(y - centerY, 2))
        H = np.exp(-1 * np.power(D, 2) / (2 * np.power(D0, 2)))
        fft = np.fft.fft2(img)  # 傅里叶变换
        fft_shift = np.fft.fftshift(fft)  # 中心化
        fft_mask = fft_shift * H
        ifft_shift = np.fft.ifftshift(fft_mask)
        ifft = np.fft.ifft2(ifft_shift)
        ifft_abs = np.abs(ifft)
        return cv2.convertScaleAbs(ifft_abs)

    # 形态处理
    @Slot(int, int)
    @process()
    def erode(self, shape, size, src):
        # 读取原始图像
        img = src
        kernel = cv2.getStructuringElement(shape, (size, size))
        imgErode = cv2.erode(img, kernel=kernel)  # 图像腐蚀
        return imgErode

    @Slot(int, int)
    @process()
    def dilate(self, shape, size, src):
        # 读取原始图像
        img = src
        kernel = cv2.getStructuringElement(shape, (size, size))
        imgDilate = cv2.dilate(img, kernel=kernel)  # 图像膨胀
        return imgDilate

    @Slot(int, int)
    @process()
    def open(self, shape, size, src):
        # 读取原始图像
        img = src
        kernel = cv2.getStructuringElement(shape, (size, size))
        imgOpen = cv2.morphologyEx(img, cv2.MORPH_OPEN, kernel)
        return imgOpen

    @Slot(int, int)
    @process()
    def close(self, shape, size, src):
        # 读取原始图像
        img = src
        kernel = cv2.getStructuringElement(shape, (size, size))
        imgClose = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)
        return imgClose

    # 噪声添加
    @Slot(float, float)
    @process()
    def gaussNoise(self, mu, sigma, src):
        img = src
        noiseGauss = np.random.normal(mu, sigma, img.shape)
        imgGaussNoise = img + noiseGauss
        return cv2.convertScaleAbs(imgGaussNoise)

    @Slot(float)
    @process()
    def expNoise(self, a, src):
        # 指数噪声 (Exponential noise)
        img = src
        noiseExponent = np.random.exponential(scale=a, size=img.shape)
        imgExponentNoise = img + noiseExponent
        return cv2.convertScaleAbs(imgExponentNoise)

    @Slot(float, float)
    @process()
    def uniformNoise(self, mean, sigma, src):
        img = src
        a = 2 * mean - np.sqrt(12 * sigma)
        b = 2 * mean + np.sqrt(12 * sigma)
        noiseUniform = np.random.uniform(a, b, img.shape)
        imgUniformNoise = img + noiseUniform
        return cv2.convertScaleAbs(imgUniformNoise)

    @Slot(float, float)
    @process()
    def pepperSaltNoise(self, ps, pp, src):
        img = src
        mask = np.random.choice((0, 0.5, 1), size=img.shape[:2], p=[pp, (1 - ps - pp), ps])
        imgChoiceNoise = img.copy()
        imgChoiceNoise[mask == 1] = 255
        imgChoiceNoise[mask == 0] = 0
        return imgChoiceNoise

    # 风格迁移
    @Slot(str)
    @process()
    def styleTransfer(self, style, src):
        predictor = LapStylePredictor(style)
        img = predictor.run(src)
        return img


if __name__ == "__main__":
    app = QApplication(sys.argv)
    engine = QQmlApplicationEngine()
    engine.load(os.fspath(Path(__file__).resolve().parent / "main.qml"))
    if not engine.rootObjects():
        sys.exit(-1)
    sys.exit(app.exec())