import traceback

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


def get_binary_img_(img):
    gray_img = img
    if len(img.shape) > 2:
        gray_img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    binary_img = cv.Canny(gray_img, 80, 150)
    return binary_img


def get_morp_dilate_(binary_img):
    kernel = cv.getStructuringElement(cv.MORPH_RECT, (3, 3))
    # morp_dilate = cv.morphologyEx(binaryImg, cv.MORPH_DILATE, kernel=(1, 3), iterations=3)
    # morp_dilate = cv.morphologyEx(morp_dilate, cv.MORPH_DILATE, kernel=(3, 1), iterations=3)
    # morp_dilate = cv.morphologyEx(binaryImg, cv.MORPH_DILATE, kernel=(11, 11), iterations=3)
    morp_dilate = cv.morphologyEx(binary_img, cv.MORPH_DILATE, kernel=kernel, iterations=3)
    return morp_dilate


def get_water_img_(img, morp_dilate):
    # 寻找图像轮廓 返回修改后的 图像的轮廓  以及它们的层次
    # contours, hierarchy = cv.findContours(gray_img, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
    # contours, hierarchy = cv.findContours(thresh, cv.RETR_TREE, cv.CHAIN_APPROX_NONE)
    # contours, hierarchy = cv.findContours(thresh, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
    contours, hierarchy = cv.findContours(morp_dilate, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
    # 32位有符号整数类型，
    marks = np.zeros(morp_dilate.shape[:2], np.int32)
    # 绘制每一个轮廓
    for index in range(len(contours)):
        # 对marks进行标记，对不同区域的轮廓使用不同的亮度绘制，相当于设置注水点，有多少个轮廓，就有多少个轮廓
        # 图像上不同线条的灰度值是不同的，底部略暗，越往上灰度越高
        marks = cv.drawContours(marks, contours, index, (index, index, index), 1, 8, hierarchy)

    # 使用分水岭算法
    # 经过watershed函数的处理，不同区域间的值被置为-1（边界）没有标记清楚的区域被置为0，其他每个区域的值保持不变：1，2，...，contours.size()
    marks_water = cv.watershed(img, marks)
    return marks_water


def get_mask_img_(morp_dilate, file_dir):
    contours, hierarchy = cv.findContours(morp_dilate, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
    # 32位有符号整数类型，
    marks = np.zeros(morp_dilate.shape[:2], np.int32)
    for index in range(len(contours)):
        dist = cv.pointPolygonTest(contours[index], (marks.shape[0] // 2, marks.shape[1] // 2), True)
        if dist >= 0:
            marks = cv.drawContours(marks, contours, contourIdx=index, color=1, thickness=1, lineType=8,
                                    hierarchy=hierarchy)

    edges = np.zeros((marks.shape[0] + 2, marks.shape[1] + 2), np.uint8)  # 掩码，长短需要加2个像素
    try:
        cv.floodFill(marks, edges, (marks.shape[0] // 2, marks.shape[1] // 2), 1, cv.FLOODFILL_MASK_ONLY)  # 漫水填充
    except Exception as e:
        if file_dir:
            print(file_dir)
            print(e)
            print("=================")
            print(traceback.format_exc())
        # raise e
        marks = np.ones(morp_dilate.shape[:2], np.int32)
    return marks


def get_binary_img(binary_img, mask):
    masked_binary_img = cv.bitwise_and(binary_img, binary_img, mask=mask.astype('uint8'))
    return masked_binary_img


def get_water_img(img, morp_dilate, mask):
    water_img = get_water_img_(img, morp_dilate)
    masked_water = cv.bitwise_and(water_img, water_img, mask=mask.astype('uint8'))
    return masked_water


def get_more_dim(img, file_dir, source_img=None):
    if source_img is None:
        source_img = img
    # img: ndarray: 852, 847, 3
    binary_img = get_binary_img_(img)
    morp_dilate = get_morp_dilate_(binary_img)
    mask = get_mask_img_(morp_dilate, file_dir)

    masked_binary_img = get_binary_img(binary_img, mask)
    masked_water = get_water_img(source_img, morp_dilate, mask)
    # print(f"masked_binary_img shape:{masked_binary_img.shape} masked_water shape:{masked_water.shape}")
    # print(f"type(masked_binary_img):{type(masked_binary_img)} type(masked_water):{type(masked_water)}")
    # return np.stack((masked_binary_img, mask), axis=0)
    return masked_binary_img, masked_water