import gi
gi.require_version('Gst', '1.0')
from gi.repository import Gst
import os
import numpy as np
import cv2
import hailo

from threading import Thread, Lock
import queue, time
from luma.core.interface.serial import i2c
from luma.oled.device import sh1106
from PIL import Image, ImageDraw, ImageFont
from adafruit_servokit import ServoKit
from datetime import datetime

from hailo_apps_infra.hailo_rpi_common import (
    get_caps_from_pad,
    get_numpy_from_buffer,
    app_callback_class,
)
from fdetection_pipeline import GStreamerDetectionApp

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

oled_queue = queue.Queue()
oled_lock = Lock()

def oled_worker():
    while True:
        text_lines = oled_queue.get()
        with oled_lock:
            image = Image.new("1", (device.width, device.height))
            draw = ImageDraw.Draw(image)
            y = 0
            for line in text_lines:
                draw.text((0, y), line, font=font, fill=255)
                y += 20
            device.display(image)
        oled_queue.task_done()

Thread(target=oled_worker, daemon=True).start()

# ----------------------------
# SERVO INIT
# ----------------------------
kit = ServoKit(channels=16, address=0x40)
PAN_MIN, PAN_MAX = 0, 180
TILT_MIN, TILT_MAX = 10, 180
pan_angle, tilt_angle = 90, 145
kit.servo[0].angle = pan_angle
kit.servo[1].angle = tilt_angle

# Tracking tuning
PAN_GAIN = 0.004       # reduce gain (was 0.015)
TILT_GAIN = 0.008      # reduce gain (was 0.015)
MAX_DELTA = 2          # smaller step per frame (was 5)
LOST_THRESHOLD = 10    # wait longer before declaring face lost
SUSTAIN_THRESHOLD = 5  # require more stable frames before moving
prev_delta = 0
state = "SEARCH"
lost_frames = 0
sustain_frames = 0
off_center_frames = 0
frame_count = 0

# Globals for stabilized face tracking
current_face = None
current_area = 0
switch_counter = 0
SWITCH_THRESHOLD = 5   # require 5 consecutive frames before switching target

# For SEARCH sweep when no faces detected
sweep_dx = 1       # direction for PAN: +1 means increasing angle
sweep_dy = 1       # direction for TILT
PAN_STEP_SEARCH = 2
TILT_STEP_SEARCH = 2


# ----------------------------
# Utility
# ----------------------------
def select_face(face_boxes):
    """Pick a stable face target, not just the max-area every frame."""
    global current_face, current_area, switch_counter

    if not face_boxes:
        current_face, current_area, switch_counter = None, 0, 0
        return None

    # Find biggest face in this frame
    max_box, max_area = None, 0
    for fbox in face_boxes:
        area = (fbox[2]-fbox[0]) * (fbox[3]-fbox[1])
        if area > max_area:
            max_area = area
            max_box = fbox

    if current_face is None:
        # No locked face yet → take max
        current_face, current_area = max_box, max_area
        return current_face

    # If same face (IoU overlap), keep it
    iou_score = iou(current_face, max_box)
    if iou_score > 0.3:  # still overlapping → stick to it
        current_face, current_area = max_box, max_area
        switch_counter = 0
        return current_face

    # If new face is bigger but not stable yet
    if max_area > current_area * 1.2:  # 20% bigger
        switch_counter += 1
        if switch_counter >= SWITCH_THRESHOLD:
            # Confirm switch
            current_face, current_area = max_box, max_area
            switch_counter = 0
    else:
        switch_counter = 0

    return current_face

def iou(boxA, boxB):
    xA1, yA1, xA2, yA2 = boxA
    xB1, yB1, xB2, yB2 = boxB
    x_left, y_top = max(xA1, xB1), max(yA1, yB1)
    x_right, y_bottom = min(xA2, xB2), min(yA2, yB2)
    if x_right < x_left or y_bottom < y_top:
        return 0.0
    inter_area = (x_right - x_left) * (y_bottom - y_top)
    boxA_area = (xA2 - xA1) * (yA2 - yA1)
    boxB_area = (xB2 - xB1) * (yB2 - yB1)
    return inter_area / float(boxA_area + boxB_area - inter_area)


# ----------------------------
# GStreamer callback with tracking
# ----------------------------
def app_callback(pad, info, user_data):
    global pan_angle, tilt_angle, state, sustain_frames, lost_frames
    global frame_count
    frame_count +=1
    buffer = info.get_buffer()
    if buffer is None:
        return Gst.PadProbeReturn.OK

    format, width, height = get_caps_from_pad(pad)
    if not (format and width and height):
        return Gst.PadProbeReturn.OK

    success, mapinfo = buffer.map(Gst.MapFlags.READ)
    if not success:
        return Gst.PadProbeReturn.OK

    frame = np.frombuffer(mapinfo.data, dtype=np.uint8).reshape((height, width, 3))
    buffer.unmap(mapinfo)

    roi = hailo.get_roi_from_buffer(buffer)
    detections = roi.get_objects_typed(hailo.HAILO_DETECTION)

    face_boxes, person_boxes = [], []
    for det in detections:
        label = det.get_label()
        bbox = det.get_bbox()
        x1 = int(bbox.xmin() * width)
        y1 = int(bbox.ymin() * height)
        x2 = int(bbox.xmax() * width)
        y2 = int(bbox.ymax() * height)
        if label == "face":
            face_boxes.append((x1, y1, x2, y2))
        elif label == "person":
            person_boxes.append((x1, y1, x2, y2))

    # pick stable face instead of instant max
    stable_box = select_face(face_boxes)
    face_center = None
    if stable_box:
        face_center = ((stable_box[0] + stable_box[2]) // 2,
                       (stable_box[1] + stable_box[3]) // 2)

    stable_face_count = len(face_boxes)

    # -------------------- 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 = width//2, height//2
                error_x, error_y = cx - frame_cx, cy - frame_cy

                # ---------------- PAN (smoothed + hysteresis) ----------------
                global prev_delta, off_center_frames
                enter_deadzone = 12
                exit_deadzone = 6

                if abs(error_x) < exit_deadzone:
                         smoothed_delta = 0
                         off_center_frames = 0
                elif abs(error_x) > enter_deadzone:
                         off_center_frames += 1
                         if off_center_frames >= 3:  # require 3 frames to confirm drift
                                 smoothed_delta = 0.2 * error_x + 0.8 * prev_delta
                                 prev_delta = smoothed_delta
                                 smoothed_delta = max(min(smoothed_delta, 2), -2)
                         else:
                             smoothed_delta = 0
                else:
                     smoothed_delta = 0

		# Apply only if movement is needed
                if smoothed_delta != 0 and frame_count % 3 == 0:  # throttle to ~10Hz
                         pan_angle = max(PAN_MIN, min(PAN_MAX, pan_angle + smoothed_delta))

                # ---------------- TILT (unchanged) ----------------
                target_tilt = tilt_angle + error_y * TILT_GAIN
                tilt_delta = max(-MAX_DELTA, min(MAX_DELTA, target_tilt - tilt_angle))
                tilt_angle = max(TILT_MIN, min(TILT_MAX, tilt_angle + tilt_delta))

                # ---------------- Apply to servos ----------------
                kit.servo[0].angle = pan_angle
                kit.servo[1].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")

    elif state == "SEARCH":
        global sweep_dx, sweep_dy

        if stable_face_count > 0 and face_center:
            # Face appeared → switch to TRACK
            state = "TRACK"
            sustain_frames = 1
            lost_frames = 0
            print(f"[{datetime.now().strftime('%H:%M:%S')}] Face found during SEARCH → TRACK")
        else:
            # No faces → move diagonally across frame
            pan_angle += sweep_dx * PAN_STEP_SEARCH
            tilt_angle += sweep_dy * TILT_STEP_SEARCH

            # Bounce PAN at limits
            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 at limits
            if tilt_angle >= TILT_MAX:
                tilt_angle = TILT_MAX
                sweep_dy = -1
            elif tilt_angle <= TILT_MIN:
                tilt_angle = TILT_MIN
                sweep_dy = 1

            # Apply to servos
            kit.servo[0].angle = pan_angle
            kit.servo[1].angle = tilt_angle
    # -------------------- OLED --------------------
    now_str = datetime.now().strftime("%H:%M:%S")
    # Print face count to terminal
    raw_face_count = len(face_boxes)
    if stable_box:
        stable_face_count = 1
    else:
        stable_face_count = 0

    # Terminal
    #print(f"[{now_str}] Raw Faces: {raw_face_count}, Stable: {stable_face_count}, Pan: {int(pan_angle)}, Tilt: {int(tilt_angle)}")

    # OLED
    oled_queue.put_nowait([
        f"{now_str}",
        f"Faces:{raw_face_count} / {stable_face_count}",
        f"Pan:{int(pan_angle)} Tilt:{int(tilt_angle)}"
    ])


    return Gst.PadProbeReturn.OK

# ----------------------------
# Main Runner
# ----------------------------
if __name__ == "__main__":
    class UserCallback(app_callback_class):
        def __init__(self):
            super().__init__()

    Gst.init(None)
    user_data = UserCallback()
    app = GStreamerDetectionApp(app_callback, user_data)
    app.run()
