from ultralytics import YOLO
import cv2
import subprocess
import numpy as np
from luma.core.interface.serial import i2c
from luma.oled.device import sh1106
from PIL import Image, ImageDraw, ImageFont
from adafruit_servokit import ServoKit
import time
from datetime import datetime
from collections import deque

# -------- OLED SETUP --------
serial = i2c(port=1, address=0x3C)
device = sh1106(serial)
device.clear()
font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 14)

# -------- SERVO SETUP --------
#i2cdetect -y 1
kit = ServoKit(channels=16, address=0x40)

# Servo limits
PAN_MIN, PAN_MAX = 0, 180
TILT_MIN, TILT_MAX = 10, 180

# Start angles (safe defaults)
pan_angle = 90
tilt_angle = 145
kit.servo[0].angle = pan_angle
kit.servo[1].angle = tilt_angle
time.sleep(0.5)

# History for stable counting
face_history = deque(maxlen=5)   # smooth over last 5 frames
stable_face_count = 0

# --- Smoothly move to center (90, 90) ---
target_pan, target_tilt = 90, 90
while abs(pan_angle - target_pan) > 1 or abs(tilt_angle - target_tilt) > 1:
    if pan_angle < target_pan:
        pan_angle += 1
    elif pan_angle > target_pan:
        pan_angle -= 1

    if tilt_angle < target_tilt:
        tilt_angle += 1
    elif tilt_angle > target_tilt:
        tilt_angle -= 1

    kit.servo[0].angle = pan_angle
    kit.servo[1].angle = tilt_angle
    time.sleep(0.002)

pan_angle, tilt_angle = target_pan, target_tilt

# -------- YOLO MODEL --------
model = YOLO("yolov8n-face-lindevs.pt")

# Camera setup (rotate at capture so frames are upright)
command = [
    "rpicam-vid", "-t", "0", "--inline", "-o", "-", "--codec", "yuv420",
    "--width", "640", "--height", "480", "--framerate", "30" 
]
command1 = [
    "rpicam-vid", "-t", "0", "--inline", "-o", "-", "--codec", "yuv420",
    "--width", "640", "--height", "480", "--framerate", "30",
    "--rotation", "180"
]     #to invert the camera use this command
proc = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, bufsize=10**8)
frame_size = (640, 480)
yuv_size = int(frame_size[0] * frame_size[1] * 1.5)

# -------- TRACKING VARS & BEHAVIOR TUNING --------
max_faces = 0
best_pan, best_tilt = pan_angle, tilt_angle

# Search sweep diagonal
sweep_dx, sweep_dy = 1, 1
PAN_STEP_SEARCH = 1
TILT_STEP_SEARCH = 1

# thresholds & smoothing
LOST_THRESHOLD = 8          # frames without faces before switching to SEARCH
SUSTAIN_THRESHOLD = 3       # frames with faces to commit to TRACK
PAN_GAIN = 0.015            # proportional gain for pan (smaller = smoother)
TILT_GAIN = 0.015           # proportional gain for tilt
MAX_DELTA = 5               # max degrees change per loop to avoid jitter

# After full scan found best -> go to it and hold
HOLD_SECONDS = 6

# State machine: 'SEARCH', 'TRACK', 'GOTO_BEST'
state = "SEARCH"
lost_frames = 0
sustain_frames = 0
goto_best_start_time = None

try:
    while True:
        raw = proc.stdout.read(yuv_size)
        if len(raw) < yuv_size:
            break

        # Convert YUV420 → BGR
        yuv = np.frombuffer(raw, dtype=np.uint8).reshape((int(frame_size[1]*1.5), frame_size[0]))
        bgr = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR_I420)

        # Run YOLO
        results = model.predict(source=bgr, verbose=False)

        face_count = 0
        face_center = None
        max_area = 0

        for r in results:
            for box in r.boxes:
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                face_count += 1
                area = (x2 - x1) * (y2 - y1)
                if area > max_area:
                    max_area = area
                    face_center = ((x1 + x2) // 2, (y1 + y2) // 2)

                # Draw rectangle and label on image
                cv2.rectangle(bgr, (x1, y1), (x2, y2), (0, 255, 0), 2)
                cv2.putText(bgr, "Face", (x1, y1 - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

        # update stable count history
        face_history.append(face_count)
        stable_face_count = max(set(face_history), key=face_history.count)

        ### NEW: Save snapshot if stable faces detected
        if stable_face_count > 0:

            # Format date-time string
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            # Overlay the text on the image (bottom-left corner)
            cv2.putText(bgr,timestamp,(10, bgr.shape[0] - 10),   # x=10, y=height-10
                        cv2.FONT_HERSHEY_SIMPLEX,0.8,                       # font scale
                        (0, 255, 0),               # green text
                        2                          # thickness
            )
            cv2.imwrite("output.jpg", bgr)
            #print(f"[{datetime.now().strftime('%H:%M:%S')}] Snapshot saved: output.jpg")

        # -------------------------------------------
        # State: TRACK (we have a face and are actively tracking)
        # -------------------------------------------
        if state == "TRACK":
            if stable_face_count > 0 and face_center:
                sustain_frames += 1
                lost_frames = 0

                if sustain_frames >= SUSTAIN_THRESHOLD:
                    cx, cy = face_center
                    frame_cx, frame_cy = frame_size[0] // 2, frame_size[1] // 2
                    error_x = cx - frame_cx
                    error_y = cy - frame_cy

                    target_pan = pan_angle - error_x * PAN_GAIN
                    target_tilt = tilt_angle + error_y * TILT_GAIN

                    pan_delta = max(-MAX_DELTA, min(MAX_DELTA, target_pan - pan_angle))
                    tilt_delta = max(-MAX_DELTA, min(MAX_DELTA, target_tilt - tilt_angle))

                    pan_angle += pan_delta
                    tilt_angle += tilt_delta

                    pan_angle = max(PAN_MIN, min(PAN_MAX, pan_angle))
                    tilt_angle = max(TILT_MIN, min(TILT_MAX, tilt_angle))

                    kit.servo[0].angle = pan_angle
                    kit.servo[1].angle = tilt_angle

                    if stable_face_count >= max_faces and sustain_frames >= SUSTAIN_THRESHOLD:
                        if stable_face_count > max_faces:
                            max_faces = stable_face_count
                            best_pan, best_tilt = pan_angle, tilt_angle

            else:
                sustain_frames = 0
                lost_frames += 1
                if lost_frames >= LOST_THRESHOLD:
                    state = "SEARCH"
                    print(f"[{datetime.now().strftime('%H:%M:%S')}] Lost faces → SEARCH")

        # -------------------------------------------
        # State: SEARCH (diagonal criss-cross)
        # -------------------------------------------
        elif state == "SEARCH":
            if stable_face_count > 0 and face_center:
                state = "TRACK"
                sustain_frames = 1
                lost_frames = 0
                print(f"[{datetime.now().strftime('%H:%M:%S')}] Face found during SEARCH → TRACK")
                continue

            # move diagonally
            #pan_angle += sweep_dx * PAN_STEP_SEARCH
            #tilt_angle += sweep_dy * TILT_STEP_SEARCH

            # adaptive search speed
            if lost_frames > 30:   # after ~1 sec of no faces
                pan_step = 3
                tilt_step = 2
            else:
                pan_step = 1
                tilt_step = 1

            pan_angle += sweep_dx * pan_step
            tilt_angle += sweep_dy * tilt_step

            # bounce pan
            if pan_angle >= PAN_MAX:
                pan_angle = PAN_MAX
                sweep_dx = -1
            elif pan_angle <= PAN_MIN:
                pan_angle = PAN_MIN
                sweep_dx = 1

            # bounce tilt
            if tilt_angle >= TILT_MAX:
                tilt_angle = TILT_MAX
                sweep_dy = -1
            elif tilt_angle <= TILT_MIN:
                tilt_angle = TILT_MIN
                sweep_dy = 1

            kit.servo[0].angle = pan_angle
            kit.servo[1].angle = tilt_angle
            #time.sleep(0.05)

            if stable_face_count > 0:
                sustain_frames += 1
                if sustain_frames >= SUSTAIN_THRESHOLD and stable_face_count >= max_faces:
                    max_faces = stable_face_count
                    best_pan, best_tilt = pan_angle, tilt_angle
            else:
                sustain_frames = 0

        # -------------------------------------------
        # State: GOTO_BEST
        # -------------------------------------------
        elif state == "GOTO_BEST":
            pan_step = max(-MAX_DELTA, min(MAX_DELTA, best_pan - pan_angle))
            tilt_step = max(-MAX_DELTA, min(MAX_DELTA, best_tilt - tilt_angle))
            if abs(pan_step) < 0.5 and abs(tilt_step) < 0.5:
                kit.servo[0].angle = best_pan
                kit.servo[1].angle = best_tilt
                if goto_best_start_time is None:
                    goto_best_start_time = time.time()
                if stable_face_count > 0:
                    state = "TRACK"
                    sustain_frames = 1
                    lost_frames = 0
                    print(f"[{datetime.now().strftime('%H:%M:%S')}] Face reappeared at best → TRACK")
                elif time.time() - goto_best_start_time >= HOLD_SECONDS:
                    print(f"[{datetime.now().strftime('%H:%M:%S')}] Hold done; resuming SEARCH")
                    max_faces = 0
                    best_pan, best_tilt = pan_angle, tilt_angle
                    state = "SEARCH"
            else:
                pan_angle += pan_step
                tilt_angle += tilt_step
                kit.servo[0].angle = pan_angle
                kit.servo[1].angle = tilt_angle
                time.sleep(0.02)

        # --- OLED update & logging ---
        image = Image.new("1", (device.width, device.height))
        draw = ImageDraw.Draw(image)
        now_str = datetime.now().strftime("%H:%M:%S")
        draw.text((10, 0), f"{now_str}", font=font, fill=255)
        draw.text((10, 20), f"Face={stable_face_count}:{max_faces}", font=font, fill=255)
        draw.text((10, 40), f"Pn:{int(pan_angle)} Tlt:{int(tilt_angle)}", font=font, fill=255)
        device.display(image)

        #print(f"{now_str}  STATE={state}  Faces={stable_face_count}  Max={max_faces}  Pan={int(pan_angle)} Tilt={int(tilt_angle)}")

except KeyboardInterrupt:
    print("Stopped by user")

finally:
    proc.terminate()
    proc.wait()
    device.clear()
    print("Camera closed, OLED cleared, Servos stopped")


