import cv2
import numpy as np
import RPi.GPIO as GPIO
import time

# ====================
# GPIO 세팅
# ====================
MOTOR_PIN = 18
GPIO.setmode(GPIO.BCM)
GPIO.setup(MOTOR_PIN, GPIO.OUT)

def move_car():
    GPIO.output(MOTOR_PIN, GPIO.HIGH)

def stop_car():
    GPIO.output(MOTOR_PIN, GPIO.LOW)

# ====================
# HSV 색상 범위 정의
# ====================
HSV_RANGES = {
    'red': [
        (np.array([0, 100, 100]), np.array([10, 255, 255])),
        (np.array([160, 100, 100]), np.array([179, 255, 255]))
    ],
    'green': [
        (np.array([40, 50, 50]), np.array([90, 255, 255]))
    ]
}

def get_color_mask(hsv_img, color_name):
    mask = None
    for lower, upper in HSV_RANGES[color_name]:
        current_mask = cv2.inRange(hsv_img, lower, upper)
        if mask is None:
            mask = current_mask
        else:
            mask = cv2.bitwise_or(mask, current_mask)
    return mask

# ====================
# 카메라 및 인식 루프
# ====================
cap = cv2.VideoCapture(0)
if not cap.isOpened():
    print("Camera not found.")
    exit()

print("Starting traffic light detection...")

try:
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        frame = cv2.resize(frame, (640, 480))
        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

        red_mask = get_color_mask(hsv, 'red')
        green_mask = get_color_mask(hsv, 'green')

        red_pixels = cv2.countNonZero(red_mask)
        green_pixels = cv2.countNonZero(green_mask)

        # 기준값 이상이면 신호등이 있다고 판단
        if red_pixels > 1000:
            print("🟥 Red Light Detected – STOP")
            stop_car()
        elif green_pixels > 1000:
            print("🟩 Green Light Detected – GO")
            move_car()
        else:
            print("⚪ No clear signal – STOP (Safe)")
            stop_car()

        # 디버깅용 화면 보기
        cv2.imshow("Frame", frame)
        if cv2.waitKey(1) & 0xFF == ord("q"):
            break

finally:
    cap.release()
    cv2.destroyAllWindows()
    GPIO.cleanup()