# import cv2
# import numpy as np
# from scipy.spatial import distance as dist, distance
#
# class Robot:
#     def __init__(self, pos=None, team='-no team!-', ID='-no ID!-'):
#         self.pos = pos if pos is not None else []
#         self.team = team
#         self.ID = ID
#         self.circles = []
#
#     def add_marking(self, circle=None):
#         if circle is None:
#             circle = [0, 0, [0, 0, 0]]
#         self.circles.append(circle)
#
# class Ball:
#     def __init__(self, pos=None):
#         self.pos = pos if pos is not None else []
#
# ball = Ball()
# robotList = []
# robotMarks = []
#
# def Color_Detection(blue, green, red):
#     if blue >= 220 and green <= 50 and red <= 50:
#         return 'Blue'
#     if blue <= 50 and green >= 200 and red >= 200:
#         return 'Yellow'
#     if blue >= 200 and green <= 50 and red >= 200:
#         return 'Purple'
#     if blue <= 50 and green >= 220 and red <= 50:
#         return 'Green'
#     if blue <= 50 and green <= 200 and red >= 180:
#         return 'Orange'
#     return 'Unidentified'
#
#
# def IdentifyCircles(img, circle):
#     global ball
#
#     x, y = int(circle[0]), int(circle[1])
#     blue, green, red = img[y, x, 0], img[y, x, 1], img[y, x, 2]
#     color = Color_Detection(blue, green, red)
#
#     if color == 'Blue' or color == 'Yellow':
#         robotList.append(Robot([x, y], color))
#     elif color == 'Green' or color == 'Purple':
#         robotMarks.append([x, y, color])
#         print('ROBOT FOUND')
#     elif color == 'Orange':
#         ball.pos = [x, y]
#         print(f"Ball found at ({x}, {y})")
#
#
# def assignIDmarks():
#     if robotList is not None:
#         for idx, robot in enumerate(robotList):
#             distances = []
#
#             for i, mark in enumerate(robotMarks):
#                 mark_dist = distance.euclidean(mark[:2], robot.pos)
#                 distances.append((i, mark_dist))
#             distances.sort(key=lambda x: x[1])
#             closest_marks_indices = [i for i, _ in distances[:4]]
#             robot.circles = [robotMarks[i] for i in closest_marks_indices]
#             robot.ID = idx + 1
#
#
# def detect_circles(image):
#     gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
#     blurred = cv2.GaussianBlur(gray, (9, 9), 0)
#     circles = cv2.HoughCircles(blurred, cv2.HOUGH_GRADIENT, 1, minDist=20, param1=50, param2=14, minRadius=15,
#                                maxRadius=50)
#     return circles
#
#
# def annotate_image(img):
#     for robot in robotList:
#         team_color = "B" if robot.team == 'Blue' else "Y"
#         sting = f'Team: {team_color} | ID: {robot.ID} | POS: {robot.pos}'
#         cv2.putText(img, sting, (robot.pos[0] + 20, robot.pos[1] - 40), cv2.FONT_HERSHEY_SIMPLEX, .75,
#                     (255, 255, 255), 2, cv2.LINE_AA)
#
#     if ball.pos:
#         cv2.putText(img, f'Ball {ball.pos}', (ball.pos[0] + 20, ball.pos[1] + 20), cv2.FONT_HERSHEY_SIMPLEX, 1,
#                     (255, 255, 255), 2, cv2.LINE_AA)
#
# # Main function
# def main():
#     global robotList, robotMarks
#     global ball
#
#     while True:
#         # Initialize globals
#         robotList = []
#         robotMarks = []
#         ball = Ball()  # Ensure ball is always an instance of Ball
#
#         # Load and process the image
#         imgpath = "/Users/mannpatel/Desktop/Project/Computer Vision/Test2.jpeg"
#         img = cv2.imread(imgpath)
#         if img is None:
#             print(f"Failed to load image at path: {imgpath}")
#             return
#
#         cv2.imshow("Original Image", img)
#
#         # Detect circles in the image
#         circles = detect_circles(img)
#
#         if circles is not None:
#             circles = np.uint16(np.around(circles))
#             for circle in circles[0, :]:
#                 IdentifyCircles(img, circle)
#                 cv2.circle(img, (circle[0], circle[1]), circle[2], (0, 255, 0), 2)
#                 cv2.circle(img, (circle[0], circle[1]), 2, (0, 0, 255), 3)
#             assignIDmarks()
#
#             for robot in robotList:
#                 print(f'There is a {robot.team} robot with these ID circles:')
#                 for mark in robot.circles:
#                     print(mark)
#
#             if ball.pos:
#                 print(f'Ball found at {ball.pos}')
#
#             for robot in robotList:
#                 if robot.pos:
#                     cv2.circle(img, (robot.pos[0], robot.pos[1]), 10, (0, 0, 0), 5)
#                     for mark in robot.circles:
#                         cv2.circle(img, (mark[0], mark[1]), 10, (0, 0, 0), 5)
#
#         else:
#             print("No circles detected")
#
#         annotate_image(img)
#         cv2.imshow("Annotated Image", img)
#         cv2.waitKey(0)
#         cv2.destroyAllWindows()
#
#
# if __name__ == "__main__":
#     main()
import cv2
import numpy as np
from scipy.spatial import distance as dist, distance


class Robot:
    def __init__(self, pos=None, team='-no team!-', ID='-no ID!-'):
        self.pos = pos if pos is not None else []
        self.team = team
        self.ID = ID
        self.circles = []

    def add_marking(self, circle=None):
        if circle is None:
            circle = [0, 0, [0, 0, 0]]
        self.circles.append(circle)


class Ball:
    def __init__(self, pos=None):
        self.pos = pos if pos is not None else []


# Initialize the ball with default position
ball = Ball()
# Initialize empty lists for robots and ID markings
robotList = []
robotMarks = []

def Color_Detection(blue, green, red):
    if blue >= 220 and green <= 50 and red <= 50:
        return 'Blue'
    if blue <= 50 and green >= 200 and red >= 200:
        return 'Yellow'
    if blue >= 200 and green <= 50 and red >= 200:
        return 'Purple'
    if blue <= 50 and green >= 220 and red <= 50:
        return 'Green'
    if blue <= 50 and green <= 200 and red >= 220:
        return 'Orange'
    return 'Unidentified'

# def Color_Detection(blue, green, red):
#     if blue == 246 and green == 0 and red == 0:
#         return 'Blue'
#     if blue <= 0 and green >= 250 and red > 350:
#         return 'Yellow'
#     if blue == 247 and green == 51 and red == 235:
#         return 'Purple'
#     if blue == 77 and green == 252 and red == 118:
#         return 'Green'
#     if blue == 50 and green == 113 and red == 228:
#         return 'Orange'
#     return 'Unidentified'


def IdentifyCircles(img, circle):
    global ball

    x, y = int(circle[0]), int(circle[1])
    blue, green, red = img[y, x, 0], img[y, x, 1], img[y, x, 2]
    color = Color_Detection(blue, green, red)

    # Debugging statements
    print(f"Circle at ({x}, {y}) with BGR ({blue}, {green}, {red}) detected as {color}")

    if color == 'Blue' or color == 'Yellow':
        robotList.append(Robot([x, y], color))
    elif color == 'Green' or color == 'Purple':
        robotMarks.append([x, y, color])
        print('ROBOT FOUND')
    elif color == 'Orange':
        ball.pos = [x, y]
        print(f"Ball found at ({x}, {y})")


def assignIDmarks():
    if robotList is not None:
        for idx, robot in enumerate(robotList):
            distances = []

            for i, mark in enumerate(robotMarks):
                mark_dist = distance.euclidean(mark[:2], robot.pos)
                distances.append((i, mark_dist))
            distances.sort(key=lambda x: x[1])
            closest_marks_indices = [i for i, _ in distances[:4]]
            robot.circles = [robotMarks[i] for i in closest_marks_indices]
            robot.ID = idx + 1


def detect_circles(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (9, 9), 0)
    circles = cv2.HoughCircles(blurred, cv2.HOUGH_GRADIENT, 1, minDist=20, param1=50, param2=14, minRadius=15,
                               maxRadius=50)
    return circles


def annotate_image(img):
    for robot in robotList:
        team_color = "B" if robot.team == 'Blue' else "Y"
        cv2.putText(img, f'{team_color}', (robot.pos[0] + 20, robot.pos[1] - 40), cv2.FONT_HERSHEY_SIMPLEX, .75,
                    (255, 255, 255), 2, cv2.LINE_AA)
        cv2.putText(img, f'ID{robot.ID}', (robot.pos[0] + 20, robot.pos[1] - 20), cv2.FONT_HERSHEY_SIMPLEX, .75,
                    (255, 255, 255), 2, cv2.LINE_AA)
        cv2.putText(img, f'{robot.pos}', (robot.pos[0] + 20, robot.pos[1]), cv2.FONT_HERSHEY_SIMPLEX, .75,
                    (255, 255, 255), 2, cv2.LINE_AA)

    if ball.pos:
        cv2.putText(img, f'Ball {ball.pos}', (ball.pos[0] + 20, ball.pos[1] + 20), cv2.FONT_HERSHEY_SIMPLEX, 1,
                    (255, 255, 255), 2, cv2.LINE_AA)
        cv2.circle(img, (ball.pos[0], ball.pos[1]), 10, (0, 165, 255), -1)  # Orange color for the ball


# Main function
def main():
    global robotList, robotMarks
    global ball

    # Initialize globals
    robotList = []
    robotMarks = []
    ball = Ball()  # Ensure ball is always an instance of Ball

    # Load and process the image
    imgpath = "/Users/mannpatel/Desktop/Robocup/Computer Vision/Template1.png"
    img = cv2.imread(imgpath)
    if img is None:
        print(f"Failed to load image at path: {imgpath}")
        return

    cv2.imshow("Original Image", img)

    # Detect circles in the image
    circles = detect_circles(img)

    if circles is not None:
        circles = np.uint16(np.around(circles))
        for circle in circles[0, :]:
            IdentifyCircles(img, circle)
            cv2.circle(img, (circle[0], circle[1]), circle[2], (0, 255, 0), 2)
            cv2.circle(img, (circle[0], circle[1]), 2, (0, 0, 255), 3)

        assignIDmarks()

        for robot in robotList:
            print(f'There is a {robot.team} robot with these ID {robot.ID}')
            for mark in robot.circles:
                print(mark)

        if ball.pos:
            print(f'Ball found at {ball.pos}')

        for robot in robotList:
            if robot.pos:
                cv2.circle(img, (robot.pos[0], robot.pos[1]), 10, (0, 0, 0), 5)
                for mark in robot.circles:
                    cv2.circle(img, (mark[0], mark[1]), 10, (0, 0, 0), 5)

    else:
        print("No circles detected")

    annotate_image(img)
    cv2.imshow("Annotated Image", img)

    # Use cv2.waitKey() to display the window until a key is pressed
    while True:
        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'):
            break

    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()