"""
[단계 9] YOLO 객체 감지 + 차선 추종 통합 (최종 미니 자율주행)

학습 목표:
  - YOLO로 보행자/정지표지판/신호등 감지
  - 차선 추종 + 안전 로직 결합
  - 모델 크기 vs 처리 속도 트레이드오프 체험

조작:
  s        : 자율주행 시작/정지
  + / -    : Kp 조정
  w / x    : 속도 조정
  y / u    : YOLO 실행 주기 변경 (성능 vs 정확도)
  SPACE    : 긴급 정지
  q        : 종료

⚠️ 사전 준비
  pip install ultralytics
  첫 실행 시 yolov8n.pt(약 6MB) 자동 다운로드

⚠️ 안전 수칙
  - 보행자(인형/종이 사람)를 차선 위에 두면 차량이 정지하는지 확인
  - 정지 동작이 확실해진 후에야 더 빠른 속도로 시도
"""

import cv2
import json
import socket
import numpy as np

# ──────────────────────────────────────────
# 설정
# ──────────────────────────────────────────

# TODO 1: 본인 차량 IP로 변경
ESP32_IP = "192.168.137.74"
ESP32_PORT = 4210
STREAM_URL = f"http://{ESP32_IP}:80/stream"

# 차선 인식
WHITE_LOWER = np.array([0,   0,   220], dtype=np.uint8)
WHITE_UPPER = np.array([180, 50,  255], dtype=np.uint8)

# 제어
KP_INIT = 0.3
BASE_SPEED_INIT = 80
STEERING_MIN, STEERING_MAX, STEERING_NEUTRAL = 45, 135, 90

# YOLO COCO 클래스 ID (안전에 관련된 것만)
PERSON_ID       = 0
BICYCLE_ID      = 1
CAR_ID          = 2
MOTORCYCLE_ID   = 3
STOP_SIGN_ID    = 11
TRAFFIC_LIGHT_ID = 9


def clip(v, lo, hi):
    return max(lo, min(hi, v))


def detect_lane_center(frame):
    h, w = frame.shape[:2]
    roi = frame[h // 2:, :]
    hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
    mask = cv2.inRange(hsv, WHITE_LOWER, WHITE_UPPER)
    kernel = np.ones((3, 3), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL,
                                    cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return None
    largest = max(contours, key=cv2.contourArea)
    M = cv2.moments(largest)
    if M['m00'] == 0:
        return None
    return int(M['m10'] / M['m00'])


def calculate_steering(lane_cx, frame_center_x, kp):
    if lane_cx is None:
        return STEERING_NEUTRAL
    error = frame_center_x - lane_cx
    steering = STEERING_NEUTRAL + kp * error
    return int(clip(steering, STEERING_MIN, STEERING_MAX))


def detect_objects(model, frame):
    """
    YOLO로 객체 감지.
    반환: dict {'person': bool, 'stop_sign': bool, 'vehicle': bool, ...}
    """
    results = model(frame, verbose=False)
    detected = {
        'person': False,
        'bicycle': False,
        'vehicle': False,
        'stop_sign': False,
        'traffic_light': False,
    }

    for box in results[0].boxes:
        cls = int(box.cls)
        conf = float(box.conf)
        if conf < 0.5:
            continue

        # TODO 2: 클래스 ID에 따라 detected 딕셔너리에 표시
        # 예시:
        # if cls == PERSON_ID:
        #     detected['person'] = True
        # elif cls == BICYCLE_ID or cls == MOTORCYCLE_ID:
        #     detected['bicycle'] = True
        # elif cls == CAR_ID:
        #     detected['vehicle'] = True
        # elif cls == STOP_SIGN_ID:
        #     detected['stop_sign'] = True
        # elif cls == TRAFFIC_LIGHT_ID:
        #     detected['traffic_light'] = True
        pass

    return detected, results[0].plot()


def decide_speed(base_speed, detected):
    """객체 감지 결과로 속도 결정."""
    # TODO 3: 안전 로직 구현
    #   - 보행자/자전거/정지표지판 → speed = 0
    #   - 차량 → speed = base_speed * 0.5
    #   - 그 외 → speed = base_speed
    return base_speed  # 이 줄을 수정


def send_command(sock, speed, steering):
    cmd = {"speed": int(speed), "steering": int(steering)}
    sock.sendto(json.dumps(cmd).encode('utf-8'), (ESP32_IP, ESP32_PORT))


def main():
    print("YOLO 모델 로드 중... (첫 실행 시 약 6MB 다운로드)")
    try:
        from ultralytics import YOLO
    except ImportError:
        print("ultralytics 미설치. 다음 명령으로 설치:")
        print("  pip install ultralytics")
        return

    model = YOLO('yolov8n.pt')
    print("YOLO 모델 로드 완료")

    cap = cv2.VideoCapture(STREAM_URL)
    if not cap.isOpened():
        print("스트림 연결 실패")
        return

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

    is_running = False
    kp = KP_INIT
    base_speed = BASE_SPEED_INIT
    yolo_interval = 3  # 3프레임마다 YOLO 실행
    yolo_counter = 0
    last_detected = {'person': False, 'bicycle': False,
                     'vehicle': False, 'stop_sign': False,
                     'traffic_light': False}
    last_annotated = None

    while True:
        ret, frame = cap.read()
        if not ret:
            continue
        frame = cv2.flip(frame, 1)
        h, w = frame.shape[:2]
        frame_center_x = w // 2

        # 차선 검출
        cx = detect_lane_center(frame)
        steering = calculate_steering(cx, frame_center_x, kp)

        # YOLO (간격 두고 실행)
        yolo_counter += 1
        if yolo_counter >= yolo_interval:
            last_detected, last_annotated = detect_objects(model, frame)
            yolo_counter = 0

        # 속도 결정
        if is_running:
            speed = decide_speed(base_speed, last_detected)
        else:
            speed = 0

        # 시각화
        display = last_annotated if last_annotated is not None else frame.copy()
        if cx is not None:
            cv2.circle(display, (cx, h // 2 + 30), 8, (0, 0, 255), -1)

        flags = ", ".join([k for k, v in last_detected.items() if v]) or "Clear"
        status = "RUN" if is_running else "STOP"
        cv2.putText(display, f"[{status}] Spd={speed} Steer={steering}",
                    (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2)
        cv2.putText(display, f"Detected: {flags}",
                    (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,255), 1)
        cv2.putText(display, f"YOLO every {yolo_interval} frames",
                    (10, 75), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (200,200,0), 1)

        cv2.imshow('Mini Autonomous Drive', display)

        # 송신
        send_command(sock, speed, steering)

        # 키 처리
        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'):
            break
        elif key == ord('s'):
            is_running = not is_running
        elif key == ord(' '):
            is_running = False
            send_command(sock, 0, STEERING_NEUTRAL)
        elif key in (ord('+'), ord('=')):
            kp += 0.05
        elif key == ord('-'):
            kp = max(0.05, kp - 0.05)
        elif key == ord('w'):
            base_speed = clip(base_speed + 10, 0, 200)
        elif key == ord('x'):
            base_speed = clip(base_speed - 10, 0, 200)
        elif key == ord('y'):
            yolo_interval = max(1, yolo_interval - 1)
        elif key == ord('u'):
            yolo_interval = min(10, yolo_interval + 1)

    send_command(sock, 0, STEERING_NEUTRAL)
    sock.close()
    cap.release()
    cv2.destroyAllWindows()
    print("종료")


if __name__ == "__main__":
    main()


# ════════════════════════════════════════════════════════════
# 정답
# ════════════════════════════════════════════════════════════
"""
[TODO 2]
if cls == PERSON_ID:
    detected['person'] = True
elif cls == BICYCLE_ID or cls == MOTORCYCLE_ID:
    detected['bicycle'] = True
elif cls == CAR_ID:
    detected['vehicle'] = True
elif cls == STOP_SIGN_ID:
    detected['stop_sign'] = True
elif cls == TRAFFIC_LIGHT_ID:
    detected['traffic_light'] = True

[TODO 3]
def decide_speed(base_speed, detected):
    if detected['person'] or detected['bicycle'] or detected['stop_sign']:
        return 0
    if detected['traffic_light']:
        # COCO만으로는 색을 구분 못 하므로 안전을 위해 정지
        return 0
    if detected['vehicle']:
        return int(base_speed * 0.5)
    return base_speed
"""