import os
import sys
import traceback

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


SOURCE_IMAGE = "test.jpg"


def filter_linear(image):
    kernel = np.array([[-0.1, 0.2, -0.1], [0.2, 3.0, 0.2], [-0.1, 0.2, -0.1]])
    cv2.imshow("Исходное изображение", image)
    cv2.imshow("Результат", cv2.filter2D(image.copy(), -1, kernel))


def filter_blur(image):
    cv2.imshow("Исходное изображение", image)
    cv2.imshow("Результат", cv2.blur(image.copy(), (5, 5)))


def filter_median_blur(image):
    cv2.imshow("Исходное изображение", image)
    cv2.imshow("Результат", cv2.medianBlur(image.copy(), 5))


def filter_gauss_blur(image):
    cv2.imshow("Исходное изображение", image)
    cv2.imshow("Результат", cv2.GaussianBlur(image.copy(), (9, 9), cv2.BORDER_DEFAULT))


def filter_erode(image):
    cv2.imshow("Исходное изображение", image)
    cv2.imshow("Результат", cv2.erode(image.copy(), np.ones((11, 11))))


def filter_dilate(image):
    cv2.imshow("Исходное изображение", image)
    cv2.imshow("Результат", cv2.dilate(image.copy(), np.ones((11, 11))))


def filter_morph(image):
    # TODO переписать
    image2 = image.copy()
    image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
    cv2.imshow("Image_load", image)
    kernel = np.ones((6, 6), np.uint8)
    image3 = cv2.morphologyEx(image3, cv2.MORPH_OPEN, kernel, iterations=1)
    cv2.imshow("MORPH_OPEN", image3)
    image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
    image3 = cv2.morphologyEx(image3, cv2.MORPH_CLOSE, kernel, iterations=1)
    cv2.imshow("MORPH_CLOSE", image3)
    image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
    image3 = cv2.morphologyEx(image3, cv2.MORPH_GRADIENT, kernel, iterations=1)
    cv2.imshow("MORPH_GRADIENT", image3)
    image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
    image3 = cv2.morphologyEx(image3, cv2.MORPH_TOPHAT, kernel, iterations=1)
    cv2.imshow("MORPH_TOPHAT", image3)
    image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
    image3 = cv2.morphologyEx(image3, cv2.MORPH_BLACKHAT, kernel, iterations=1)
    cv2.imshow("MORPH_BLACKHAT", image3)
    image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)


def filter_sobel(image):
    image2 = image.copy()
    image3 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
    image3 = cv2.GaussianBlur(image3, (3, 3), 0)
    im3 = cv2.Sobel(image3, cv2.CV_64F, 1, 0, ksize=5)
    im4 = cv2.Sobel(image3, cv2.CV_64F, 0, 1, ksize=5)
    im5 = cv2.Sobel(image3, cv2.CV_8UC1, 0, 1, ksize=5)
    cv2.imshow("Image_load", image)
    cv2.imshow("Result_Image_X", im3)  # x
    cv2.imshow("Result_Image_Y", im4)  # y
    cv2.imshow("Result_Gradient", im5)  # gradient


def filter_laplacian(image):
    image2 = image.copy()
    image3 = cv2.cvtColor(image2, cv2.COLOR_RGB2GRAY)
    image3 = cv2.GaussianBlur(image3, (3, 3), 0)
    im3 = cv2.Laplacian(image3, cv2.CV_64F)
    cv2.imshow("Image_load", image)
    cv2.imshow("Result_Image", im3)


def filter_canny(image):
    image2 = image.copy()
    image3 = cv2.blur(image2, (5, 5))
    image3 = cv2.Canny(image3, 100, 100)
    cv2.imshow("Image_load", image)
    cv2.imshow("Result_Image", image3)


def filter_calc_hist(image):
    image2 = image.copy()
    color = ('b', 'g', 'r')
    cv2.imshow("Image_load", image)
    for i, col in enumerate(color):
        histr = cv2.calcHist([image2], [i], None, [256], [0, 256])
        plt.plot(histr, color=col)
        plt.xlim([0, 256])
    plt.show()
    plt.close()


def filter_equalize_hist(image):
    cv2.imshow("Исходное изображение", image)
    cv2.imshow("Результат", cv2.equalizeHist(cv2.cvtColor(image.copy(), cv2.COLOR_RGB2GRAY)))


functions_list = {
    "lf": {"func": filter_linear, "help": "Linear filter"},
    "bl": {"func": filter_blur, "help": "Blur"},
    "mb": {"func": filter_median_blur, "help": "Median blur"},
    "gb": {"func": filter_gauss_blur, "help": "GaussianBlur"},
    "er": {"func": filter_erode, "help": "Erode"},
    "di": {"func": filter_dilate, "help": "Dilate"},
    "mo": {"func": filter_morph, "help": "MORPH_OPERATIONS"},
    "so": {"func": filter_sobel, "help": "Sobel"},
    "la": {"func": filter_laplacian, "help": "Laplacian"},
    "ca": {"func": filter_canny, "help": "Canny"},
    "cl": {"func": filter_calc_hist, "help": "CalcHist"},
    "eh": {"func": filter_equalize_hist, "help": "EqualizeHist"},
    "q": {"func": lambda image: sys.exit(0), "help": "Exit"},
}


def print_help():
    print("Доступные функции:")
    for k in functions_list:
        print(f"    {k} - {functions_list[k]['help']}")


def main():
    image1 = cv2.imread(SOURCE_IMAGE)
    while True:
        try:
            cmd = input(">> ")
            if cmd == "":
                continue
            elif cmd == "?":
                print_help()
            else:
                if cmd in functions_list:
                    try:
                        functions_list[cmd]["func"](image=image1.copy())
                        cv2.waitKey(0)
                        cv2.destroyAllWindows()
                    except Exception:
                        print("Произошла ошибка при выполнении функции")
                        traceback.print_exc()
                else:
                    print("Функция не найдена! введите '?' для получения справки")

        except Exception:
            traceback.print_exc()


if __name__ == "__main__":
    main()