import cv2
import time
import serial
import numpy as np
from picamera2 import Picamera2
from collections import namedtuple

TOP_WALL = 15
BOT_WALL = 425
TOP_HBOT = TOP_WALL+85
BOT_HBOT = BOT_WALL-85
RIGHT_HBOT = 640
LEFT_HBOT = 410


class AirHockeyRobot:
    def __init__(self):
        cv2.startWindowThread()
        self.cam = Picamera2()
        config = self.cam.create_preview_configuration(
            { "size": (640, 480), "format": "RGB888" },
            raw = self.cam.sensor_modes[6]
        )
        self.cam.configure(config)
        self.cam.set_controls({'FrameRate': 80})
        self.cam.start()

        self.ser = serial.Serial('/dev/ttyACM0', 9600,  timeout=0.01)
        self.prev_time = time.time()
        self.fps = 0

        self.try_again_counter = 0
        self.target_x = RIGHT_HBOT
        self.target_y = BOT_HBOT-50


    def send_data(self, x, y):
        response = self.ser.readline().decode().strip()
        if response:
            data = f"{x},{y}\n"
            self.ser.write(data.encode())

    def detect_circles(self, frame):
        # grayscale
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # dilate to remove thin lines
        kernel = np.ones((3, 3), np.uint8)
        gray = cv2.dilate(gray, kernel, iterations=1)

        # blur to remove noise
        gray = cv2.GaussianBlur(gray, (5, 5), 2)

        circles = cv2.HoughCircles(
            gray,
            cv2.HOUGH_GRADIENT,
            dp=1.2,
            minDist=50,
            param1=250,
            param2=30,
            minRadius=20,
            maxRadius=40
        )

        filtered_circles = []
        if circles is not None:
            circles = np.uint16(np.around(circles))
            circles = sorted(circles[0, :], key=lambda c: c[2], reverse=True)

            for n, circle in enumerate(circles):
                x, y, radius = circle
                if (TOP_WALL <= y <= BOT_WALL) and (0 <= x <= RIGHT_HBOT-30):
                    filtered_circles.append(circle)

                    if n == 0 or n == 1:
                        color = (0, 255, 0)
                    else:
                        color = (0, 0, 255)
                    cv2.circle(frame, (x, y), radius, color, 3)
                    cv2.circle(frame, (x, y), 2, color, 3)

        return filtered_circles


    def update_fps(self):
        curr_time = time.time()
        self.fps = 1 / (curr_time - self.prev_time)
        self.prev_time = curr_time

    def run(self):
        while True:
            frame = self.cam.capture_array()
            #frame = frame[100-85:340+85, 0:self.def_x-35]

            circles = self.detect_circles(frame)
            if len(circles) > 1:
                puck_x, puck_y, _ = circles[1]
                puck_x = puck_x.item()
                puck_y = puck_y.item()
                paddle_x, paddle_y, _ = circles[0]
                puck_dx, puck_dy = paddle_x-puck_x, paddle_y-puck_y

                # Attack
                if puck_x >= LEFT_HBOT:
                    self.target_x = puck_x
                    self.target_y = puck_y

                # Defend
                elif puck_dx != 0 and puck_dy != 0:
                    cv2.line(frame, (paddle_x, paddle_y), (puck_x, puck_y), (255, 0, 0), 2)

                    # Check if shot on goal
                    intercept_y = int(puck_y + ((RIGHT_HBOT-puck_x) / puck_dx) * puck_dy)
                    if TOP_WALL <= intercept_y <= BOT_WALL:
                        cv2.line(frame, (puck_x, puck_y), (RIGHT_HBOT, intercept_y), (255, 0, 0), 2)

                        self.target_x = RIGHT_HBOT
                        self.target_y = np.clip(intercept_y, TOP_HBOT+50, BOT_HBOT-50)


                self.target_x = np.clip(self.target_x, LEFT_HBOT+50, RIGHT_HBOT)
                self.target_y = np.clip(self.target_y, TOP_HBOT+25, BOT_HBOT-25)
                self.send_data(self.target_x, self.target_y)
                cv2.circle(frame, (self.target_x, self.target_y), 15, (255, 0, 0), -1)

            cv2.line(frame, (0, TOP_WALL), (640, TOP_WALL), (255, 255, 255), 1)
            cv2.line(frame, (0, BOT_WALL), (640, BOT_WALL), (255, 255, 255), 1)
            cv2.line(frame, (LEFT_HBOT, 0), (LEFT_HBOT, 480), (255, 255, 255), 1)

            self.update_fps()
            cv2.putText(frame, f"FPS: {int(self.fps)}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
            cv2.imshow("Frame", frame)
            cv2.waitKey(1)


robot = AirHockeyRobot()
robot.run()