import cv2
import matplotlib.pyplot as plt


def drawSquare2D(image, x, y, size, color=(0, 0, 255), thickness=1):
    """
    Draws a square on the image
    :param image:
    :param x:
    :param y:
    :param color:
    :param thickness:
    """

    h, w, _ = image.shape
    if x - size < 0 or x + size >= w or y - size < 0 or y + size >= h:
        # print("Cannot draw square")
        return None

    # tl, tr, bl, br -> top left, top right, bottom left, bottom right
    tl = (int(x - size), int(y - size))
    tr = (int(x + size), int(y - size))
    br = (int(x + size), int(y + size))
    bl = (int(x - size), int(y + size))

    # cv2.line(image, (x,y), (x,y), color, 5)
    cv2.line(image, tl, tr, color, thickness)
    cv2.line(image, tr, br, color, thickness)
    cv2.line(image, br, bl, color, thickness)
    cv2.line(image, bl, tl, color, thickness)


def drawCross2D(image, x, y, size, color=(0, 0, 255), thickness=1):
    """
    Draws a square on the image
    :param image:
    :param x:
    :param y:
    :param color:
    :param thickness:
    """

    h, w, _ = image.shape
    if x - size < 0 or x + size > w or y - size < 0 or y + size > h:
        # print("Cannot draw square")
        return None

    # tl, tr, bl, br -> top left, top right, bottom left, bottom right
    left = (int(x - size), int(y))
    right = (int(x + size), int(y))
    top = (int(x), int(y - size))
    bottom = (int(x), int(y + size))

    cv2.line(image, left, right, color, thickness)
    cv2.line(image, top, bottom, color, thickness)


def drawDiagonalCross2D(image, x, y, size, color=(0, 0, 255), thickness=1):
    """
    Draws a square on the image
    :param image:
    :param x:
    :param y:
    :param color:
    :param thickness:
    """

    h, w, _ = image.shape
    if x - size < 0 or x + size > w or y - size < 0 or y + size > h:
        # print("Cannot draw square")
        return None

    # tl, tr, bl, br -> top left, top right, bottom left, bottom right
    top_left = (int(x - size), int(y - size))
    top_right = (int(x + size), int(y - size))
    bottom_left = (int(x - size), int(y + size))
    bottom_right = (int(x + size), int(y + size))

    cv2.line(image, top_left, bottom_right, color, thickness)
    cv2.line(image, top_right, bottom_left, color, thickness)


def drawCircle(image, x, y, size, color=(0, 0, 255), thickness=1):
    """
    Draws a square on the image
    :param image:
    :param x:
    :param y:
    :param color:
    :param thickness:
    """

    # h, w, _ = image.shape
    # if x - size < 0 or x + size >= w or y - size < 0 or y + size >= h:
    #     # print("Cannot draw square")
    #     return None

    center_coordinates = (x, y)

    cv2.circle(image, center_coordinates, size, color, thickness)
    # # tl, tr, bl, br -> top left, top right, bottom left, bottom right
    # tl = (int(x - size), int(y - size))
    # tr = (int(x + size), int(y - size))
    # br = (int(x + size), int(y + size))
    # bl = (int(x - size), int(y + size))
    #
    # # cv2.line(image, (x,y), (x,y), color, 5)
    # cv2.line(image, tl, tr, color, thickness)
    # cv2.line(image, tr, br, color, thickness)
    # cv2.line(image, br, bl, color, thickness)
    # cv2.line(image, bl, tl, color, thickness)


def draw3Dcoordinatesystem(ax, handles, xc=0, yc=0, zc=0, size=0.1, update=False):
    # vec_i=[[xc, yc, zc], [xc, yc, zc], [xc, yc, zc]]
    # vec_f=[[xc + size, yc, zc],[xc, yc + size, zc],  [xc, yc, zc + size]]
    # x_vec=[[xc, xc + size],[xc, xc],[xc, xc] ]
    # y_vec= [[yc, yc],[yc, yc + size] , [yc, yc]]
    # z_vec=[[zc, zc], [zc, zc], [zc, zc + size]]
    # colors=[(1,0,0),(0,1,0),(0,0,1)]
    # # ax.plot(vec_i, vec_f,c= colors)
    # ax.plot(x_vec, y_vec, z_vec, colors)
    if update == False:
        handle_x = ax.plot([xc, xc + size], [yc, yc], [zc, zc], color='r')
        handle_y = ax.plot([xc, xc], [yc, yc + size], [zc, zc], color='g')
        handle_z = ax.plot([xc, xc], [yc, yc], [zc, zc + size], color='b')

        return [handle_x, handle_y, handle_z]
    else:
        # print(handles[0])
        handles[0][0].set_xdata([xc, xc + size])
        handles[0][0].set_ydata([yc, yc])
        handles[0][0].set_3d_properties(zs=[zc, zc])
        handles[1][0].set_xdata([xc, xc])
        handles[1][0].set_ydata([yc, yc + size])
        handles[1][0].set_3d_properties(zs=[zc, zc])
        handles[2][0].set_xdata([xc, xc])
        handles[2][0].set_ydata([yc, yc])
        handles[2][0].set_3d_properties(zs=[zc, zc + size])
    #     ax._offsets3d([xc, xc + size], [yc, yc], [zc, zc])
    #     ax._offsets3d([xc, xc], [yc, yc + size], [zc, zc])
    #     ax._offsets3d([xc, xc], [yc, yc], [zc, zc + size])