import time
import cv2
from ultralytics import YOLO
from paho.mqtt import client as mqtt
import json
import random
import threading

# bufferless VideoCapture
class VideoCapture:
    def __init__(self, name):
        self.cap = cv2.VideoCapture(name)
        self.lock = threading.Lock()
        self.t = threading.Thread(target=self._reader)
        self.t.daemon = True
        self.t.start()

    # grab frames as soon as they are available
    def _reader(self):
        while True:
            with self.lock:
                ret = self.cap.grab()
            if not ret:
                break

    # retrieve latest frame
    def read(self):
        print("read");
        with self.lock:
            _, frame = self.cap.retrieve()
        return frame
    

# Desired frame rate (in seconds)
frame_rate = 1  # 1 / fps

# MQTT Settings
mqtt_username = "mqtt-user"
mqtt_password =  "mqtt-password"

broker = 'ip'
port = 1883
topic = "your_topic/subject"
client_id = f'publish-{random.randint(0, 1000)}'
retain = False
aantal_eieren_detected = 0

# Load the YOLOv8 model
model = YOLO("./best.pt")


#video_path = "your_path"
#cap = cv2.VideoCapture(video_path)

# Alternatively: use webcam as input
# cap = cv2.VideoCapture(0)


# Alternatively: use rtsp-stream as input
url = "rtsp://admin:password@ip:8554/Streaming/Channels/101"
cap = cv2.VideoCapture(url)


def open_stream(url):
    cap = cv2.VideoCapture(url)
    return cap

def restart_stream_input():
    cap.release()
    cap = cv2.VideoCapture(url)

mqtt_client = mqtt.Client()
mqtt_client.username_pw_set(mqtt_username, mqtt_password)
mqtt_client.connect(broker, port)
mqtt_client.loop_start()

config_payload_json = {
        "object_id": client_id,
        "name": "kwartel_eieren",
        "aantal_eieren": aantal_eieren_detected
    }
def make_payload(aantal_eieren_detected):
    config_payload_json = {
        "object_id": client_id,
        "name": "kwartel_eieren",
        "aantal_eieren": aantal_eieren_detected
    }
    config_payload = json.dumps(config_payload_json);
    return config_payload;

# Eerste keer publishen van aantal eieren (0)
mqtt_client.publish(topic, json.dumps(config_payload_json), retain=retain, qos=1)
print("Published", aantal_eieren_detected, "eieren detected")

# Variables to track time
start_time = time.time()
prev_frame_time = start_time


# Start infinite loop for when the stream stops / is too slow so end is reached and we need to restart the stream
while True:
    cap = open_stream(url);
    # Loop through the video frames
    while cap.isOpened():
        # Get current time
        current_time = time.time()
        # Calculate time elapsed since last frame
        elapsed_time = current_time - prev_frame_time
        # If elapsed time is less than the desired frame rate, continue to the next iteration
        if elapsed_time < 1.0 / frame_rate:
            ret, _ = cap.read()
            continue
        # Update previous frame time
        prev_frame_time = current_time    

        # Read a frame from the video
        success, frame = cap.read()

        if success:
            # Run YOLOv8 inference on the frame;
            results = model(frame);

            # Count the number of detections in the frame
            aantal_eieren_detected = len(results[0]);

            # Print the number of detections in the frame
            print(f"Number of detections in frame: {aantal_eieren_detected}");
            payload = make_payload(aantal_eieren_detected);

            mqtt_client.publish(topic, payload, retain=retain, qos=1);

            # SHOW BOUNDING BOXES (optional)

            # Visualize the results on the frame
            # annotated_frame = results[0].plot()
            # Display the annotated frame
            # cv2.imshow("YOLOv8 Inference", annotated_frame)

            # Break the loop if 'q' is pressed
            if cv2.waitKey(1) & 0xFF == ord("q"):
                break

        else:
            # Break the loop if the end of the video is reached
            print("End of video reached")
            break
    cap.release()
    cv2.destroyAllWindows()
    time.sleep(.2);

# Release the video capture object and close the display window
cap.release()
cv2.destroyAllWindows()