import cv2


class DrawLineWidget(object):
    def __init__(self):
        self.original_image = cv2.imread('./data/1.jpg')
        self.clone = self.original_image.copy()
        self.work = self.original_image.copy()

        cv2.namedWindow('image')
        cv2.setMouseCallback('image', self.extract_coordinates)

        # List to store start/end points
        self.image_coordinates = []

    def extract_coordinates(self, event, x, y, flags, parameters):
        # Record starting (x,y) coordinates on left mouse button click
        if event == cv2.EVENT_LBUTTONDOWN:
            if len(self.image_coordinates) == 0:
                self.image_coordinates = [(x, y)]
        elif event == cv2.EVENT_MOUSEMOVE:
            if len(self.image_coordinates) > 0:
                self.work = self.clone.copy()
                # Draw line
                cv2.line(self.work, self.image_coordinates[0], (x, y), (36, 255, 12), 2)
                cv2.imshow("image", self.work)

        # Record ending (x,y) coordinates on left mouse bottom release
        elif event == cv2.EVENT_LBUTTONUP:
            self.clone = self.work.copy()
            self.image_coordinates = []
        # Clear drawing boxes on right mouse button click
        elif event == cv2.EVENT_MBUTTONDOWN:
            self.work = self.original_image.copy()
            self.clone = self.original_image.copy()
            cv2.imshow("image", self.work)

    def show_image(self):
        return self.work


if __name__ == '__main__':
    draw_line_widget = DrawLineWidget()
    while True:
        cv2.imshow('image', draw_line_widget.show_image())
        key = cv2.waitKey(1)

        # Close program with keyboard 'q'
        if key == 27:
            cv2.destroyAllWindows()
            exit(1)