import cv2
import numpy as np


class Camera():
    def __init__(self) -> None:
        self.colors = [(0, 0, 255), (0, 255, 0), (255, 0, 0)]
        self.color_names = ["Rot", "Gruen", "Blau"]

        self.lower_red = np.array([80, 160, 150])
        self.upper_red = np.array([255, 255, 255])

        self.lower_green = np.array([40, 50, 160])
        self.upper_green = np.array([80, 255, 255])

        self.lower_blue = np.array([95, 180, 90])
        self.upper_blue = np.array([130, 255, 255])

        self.video = cv2.VideoCapture(0)
        self.image = np.ndarray([])

        self.picture_counter = 0
        self.evaluate_picture = False

    def get_frame(self) -> np.ndarray:
        try:
            _, self.image = self.video.read(0)
        except Exception as err:   
            print("Can not capture the video..\n")
            print(err)
        return self.image
    
    def set_take_picture(self, take: bool):
        self.evaluate_picture = take

    def take_picture(self) -> None:
        file_name = 'score_round'+ str(self.picture_counter) + '.png'
        cv2.imwrite(file_name, self.image)
        self.picture_counter = self.picture_counter + 1 
        self.set_take_picture(False)

    def detect_color(self, image):
        hsv_img = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        count_red, count_green, count_blue= 0,0,0
        results = []

        for i, color in enumerate(self.colors):
            if i == 0:
                lower = self.lower_red
                upper = self.upper_red
            elif i == 1:
                lower = self.lower_green
                upper = self.upper_green
            elif i == 2:
                lower = self.lower_blue
                upper = self.upper_blue

            mask = cv2.inRange(hsv_img, lower, upper)

            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            center = None
            count = 0

            for contour in contours:
                if count < 3:
                    if  100 < cv2.contourArea(contour):

                        M = cv2.moments(contour)
                        if M["m00"] > 0:
                            cX = int(M["m10"] / M["m00"])
                            cY = int(M["m01"] / M["m00"])
                            center = (cX, cY)
                            count += 1

                            x, y, w, h = cv2.boundingRect(contour)
                            cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
                            cv2.putText(image, self.color_names[i], (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)

                            if i == 0:
                                count_red += 1
                            elif i == 1:
                                count_green += 1
                            elif i == 2:
                                count_blue += 1
            results.append(center)
        current_score= {'score_red': count_red,
                        'score_green': count_green,
                        'score_blue': count_blue
                        }
        return results, image, current_score

    def determine_position(self,results, img_width):
        positions = []

        for result in results:
            if result is None:
                position = 0
            else:
                x = result[0]
                if x < img_width / 3:
                    position = 3
                elif x < 2 * img_width / 3:
                    position = 2
                else:
                    position = 1

            positions.append(position)
        return positions

    def check_correct_field(self, correct_field: int):
        pass

    def current_score(self, scores: dict):
        return scores

def main():
    my_camera = Camera()

    while True:
        frame = my_camera.get_frame()
        results, image, current_score = my_camera.detect_color(frame)   

        if my_camera.evaluate_picture:
            my_camera.take_picture()
            my_camera.current_score(current_score)

        cv2.putText(frame, f"Rot: {current_score['score_red']}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, my_camera.colors[0], 2)
        cv2.putText(frame, f"Gruen: {current_score['score_green']}", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 1, my_camera.colors[1], 2)
        cv2.putText(frame, f"Blau: {current_score['score_blue']}", (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 1, my_camera.colors[2], 2)

        img_width = frame.shape[1]
        positions = my_camera.determine_position(results, img_width)

        for i, position in enumerate(positions):
            cv2.putText(image, f"{my_camera.color_names[i]}: {position}", (10, 150 + 30 * i),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, my_camera.colors[i], 2)

        cv2.line(img=image, pt1=(img_width // 3, 0), pt2=(img_width // 3, frame.shape[0]), color=(0, 0, 0), thickness=2)
        cv2.line(img=image, pt1=(2 * img_width // 3, 0), pt2=(2 * img_width // 3, frame.shape[0]), color=(0, 0, 0), thickness=2)

        cv2.imshow("Farberkennung", frame)
        print(current_score)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    my_camera.video.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()