import gradio as gr
import cv2
import numpy as np
import json

def preprocess(img):
    # Convert to grayscale
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    cv2.imwrite("1_gray.png", img_gray)

    # Blur the image
    img_blur = cv2.GaussianBlur(img_gray, (3, 3), 1)
    cv2.imwrite("2_blur.png", img_blur)

    # Detect edges with Canny
    img_canny = cv2.Canny(img_blur, 100, 200)
    cv2.imwrite("3_canny.png", img_canny)

    # Dilate the edges
    kernel = np.ones((3, 3))
    img_dilate = cv2.dilate(img_canny, kernel, iterations=1)
    cv2.imwrite("4_dilate.png", img_dilate)

    # Erode the dilated edges
    img_erode = cv2.erode(img_dilate, kernel, iterations=0)
    cv2.imwrite("5_erode.png", img_erode)

    
    return img_erode

def find_tip(points, convex_hull):
    length = len(points)
    indices = np.setdiff1d(range(length), convex_hull)

    for i in range(2):
        j = indices[i] + 2
        if j > length - 1:
            j = length - j
        if np.all(points[j] == points[indices[i - 1] - 2]):
            return tuple(points[j])

def get_length(p1, p2):
    line_length = ((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2) ** 0.5
    return line_length

def get_angle(tip, tail):
    return (np.degrees(np.arctan2(tip[1] - tail[1], tip[0] - tail[0]))*-1)
    
def get_max_distance_point(cnt, tip):
    max_distance = 0
    max_point = None

    for [x, y] in cnt:
        distance = get_length((x, y), tip)

        if distance > max_distance:
            max_distance = distance
            max_point = (x, y)

    return max_point

def find_tail(points, tip):
    tip = np.array(tip, dtype=np.float32)
    points = np.array(points, dtype=np.int32)

    # Find the tail point based on maximum distance
    tail_point = get_max_distance_point(points, tip)

    # Ensure that the tail point is sufficiently different from the tip
    if np.linalg.norm(np.array(tail_point) - np.array(tip)) > 5:
        return tail_point

    return None

def draw_arrow(img_result, tip, tail, length, angle):
    # Draw arrow on the blank image with inverted tip and tail
    cv2.arrowedLine(img_result, tuple(tail), tuple(tip), (0, 255, 0), 3)
    """
    # Add length and angle as text next to the tip point
    text_length = f"Length: {length:.2f}"
    text_angle = f"Angle: {angle:.2f}"
    
    cv2.putText(img_result, text_length, (tip[0] + 10, tip[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
    cv2.putText(img_result, text_angle, (tip[0] + 10, tip[1] - 30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
    """
def infer(image_in):
    img = cv2.imread(image_in)

    contours, hierarchy = cv2.findContours(preprocess(img), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Create a blank image to accumulate arrows
    img_result = np.zeros_like(img)

    arrows_coordinates = []
    
    for cnt in contours:
        peri = cv2.arcLength(cnt, True)
        approx = cv2.approxPolyDP(cnt, 0.025 * peri, True)
        hull = cv2.convexHull(approx, returnPoints=False)
        sides = len(hull)

        if 6 > sides > 3 and sides + 2 == len(approx):
            arrow_tip = find_tip(approx[:,0,:], hull.squeeze())
            
            
            if arrow_tip :
                cv2.drawContours(img, [cnt], -1, (0, 255, 0), 3)
                cv2.circle(img, arrow_tip, 3, (0, 0, 255), cv2.FILLED)
                arrow_tail = find_tail(approx[:,0,:], arrow_tip)
                if arrow_tail :
                    arrows_coordinates.append([list(arrow_tail), list(arrow_tip)])
                    cv2.circle(img, arrow_tail, 3, (255, 0, 0), cv2.FILLED)
                    # Calculate length and angle
                    arrow_length = get_length(arrow_tip, arrow_tail)
                    arrow_angle = get_angle(arrow_tip, arrow_tail)
                    # Draw arrow on the same blank image
                    draw_arrow(img_result, arrow_tip, arrow_tail, arrow_length, arrow_angle)



    cv2.imwrite("Image_result.png", img)
    cv2.imwrite("Arrows_on_same_blank.png", img_result)

    print(f"arrows coordinates: {arrows_coordinates}")

    return "Image_result.png", "Arrows_on_same_blank.png", f"{arrows_coordinates}"

gr.Interface(
    fn=infer,
    inputs=gr.Image(
        sources=["upload"],
        type="filepath"
    ),
    outputs=[
        gr.Image(label="detected arrows"),
        gr.Image(label="vectors"),
        gr.Textbox(label="arrows coordinates")
        
    ],
    examples = ["example.png"]
).launch()