import matplotlib.pyplot as plt
import pydicom
import cv2
import numpy as np
from skimage.measure import find_contours


# def get_mask_contour(mask):
#     mskray = mask.astype(np.uint8)
#     edged_mask = cv2.Canny(mskray, np.min(mskray), np.max(mskray))
#     return edged_mask

# Global variables
drawing = False  # True when the mouse button is pressed
points = []  # Stores points of the ROI
mask = None  # To store the final mask

def dcm2cv (dicom_img):
    # Normalize the pixel values to the range 0-255
    image_normalized = cv2.normalize(dicom_img, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX)
    # Convert to unsigned 8-bit integer type
    image_normalized = image_normalized.astype(np.uint8)
    cv_image = cv2.merge([image_normalized, image_normalized, image_normalized])
    return cv_image
# Mouse callback function
def draw_roi(event, x, y, flags, param):
    global drawing, points, mask

    if event == cv2.EVENT_LBUTTONDOWN:  # Start drawing
        drawing = True
        points.append((x, y))

    elif event == cv2.EVENT_MOUSEMOVE and drawing:  # Add points while moving
        points.append((x, y))

    elif event == cv2.EVENT_LBUTTONUP:  # Stop drawing
        drawing = False

    elif event == cv2.EVENT_RBUTTONDOWN:  # Right-click to complete the ROI
        print('Drawing Complete')
        if len(points) > 2:
            mask = np.zeros_like(image, dtype=np.uint8)
            pts = np.array(points, np.int32).reshape((-1, 1, 2))
            cv2.fillPoly(mask, [pts], (255, 255, 255))
            contours = find_contours(mask[:, :, 0], level=0.5)
            # plt.imshow(image)
            # for cn in range(len(contours)):
            #     plt.plot(contours[cn][:, 1], contours[cn][:, 0], 'r')
            # plt.show()


# Load an image (replace this with actual MR image loading)
path = r'E:\projects\knee_seg\data_tamplet\PA0\ST0\SE4\IM40'
file_dcm = pydicom.read_file(path)
image = file_dcm.pixel_array
image = dcm2cv(image)


cv2.namedWindow("Image")
cv2.setMouseCallback("Image", draw_roi)

while True:
    temp_img = image.copy()
    print(temp_img.shape)
    if points:
        cv2.polylines(temp_img, [np.array(points)], isClosed=False, color=(255, 0, 0), thickness=2)
    cv2.imshow("Image", temp_img)

    key = cv2.waitKey(1) & 0xFF
    if key == 13 and mask is not None:  # Press 'enter' to save the mask
        cv2.imwrite("mask.png", mask)
        print("Mask saved as mask.png")
        break
    elif key == 27:  # Press 'Esc' to exit
        # print(mask)
        cv2.imwrite("mask.png", mask)
        print("Break: Mask saved  ")
        break

cv2.destroyAllWindows()