import cv2
from ultralytics import YOLO
import socket
import threading
import time
import sys

# Update this with the correct path to your model
model_path = r'C:\Users\julia\OneDrive - Hogeschool West-Vlaanderen\Project One\raspi\2023-2024-projectone-ctai-Juliane-Farmer\dog_application\model\runs\best_colab.pt'

server_address = ('192.168.168.167', 8500)  # Replace with your server's IP and port

# Global vars for use in methods/threads
client_socket = None
receive_thread = None
shutdown_flag = threading.Event()  

def setup_socket_client():
    global client_socket, receive_thread
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)  # create a socket instance
    client_socket.connect(server_address)  # connect to specified server
    print("Connected to server")

    receive_thread = threading.Thread(target=receive_messages, args=(client_socket, shutdown_flag))
    receive_thread.start()

def receive_messages(sock, shutdown_flag):
    sock.settimeout(1)  # Set a timeout on the socket so we can check shutdown_flag.is_set in the loop, instead of blocking
    try:
        while not shutdown_flag.is_set():  # as long as ctrl+c is not pressed
            try:
                data = sock.recv(1024)  # try to receive 1024 bytes of data (maximum amount; can be less)
                if not data:  # when no data is received, try again (and shutdown flag is checked again)
                    break
                print("Received from server:", data.decode())  # print the received data, or do something with it
            except socket.timeout:  # when no data comes within timeout, try again
                continue
    except Exception as e:
        if not shutdown_flag.is_set():
            print(f"Connection error: {e}")
    finally:
        sock.close()

def main():
    global client_socket, receive_thread

    setup_socket_client()

    if client_socket is None:
        print("Not connected, is server running on {}:{}?".format(server_address[0], server_address[1]))
        sys.exit()
    
    # send "hello I'm connected" message
    client_socket.sendall("Hello from AI / notebook".encode())  # send a "connected" message from client > server

    model = YOLO(model_path)
    vid = cv2.VideoCapture(0)
    prev_frame = None
    prev_detections = []
    frame_change_threshold = 50000  # Threshold for frame changes (adjust as needed)

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

            # Convert the frame to grayscale and blur it to reduce noise
            gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            gray_frame = cv2.GaussianBlur(gray_frame, (21, 21), 0)

            if prev_frame is None:
                prev_frame = gray_frame
                continue

            # Compute the absolute difference between the current frame and the previous frame
            frame_diff = cv2.absdiff(prev_frame, gray_frame)
            _, thresh = cv2.threshold(frame_diff, 25, 255, cv2.THRESH_BINARY)
            frame_diff_sum = cv2.sumElems(thresh)[0]

            # Check if the frame change is significant
            if frame_diff_sum > frame_change_threshold:
                results = model(frame)
                current_detections = []

                for result in results:
                    if result.boxes:
                        for box in result.boxes:
                            x1, y1, x2, y2 = map(int, box.xyxy[0])
                            conf = box.conf[0]
                            cls = box.cls[0]

                            # Store current detections
                            emotion = model.names[int(cls)]
                            current_detections.append((emotion, conf))

                            # Draw rectangle and label on the frame
                            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                            label = f'{emotion} {conf:.2f}'
                            cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

                # Print and send new detections if they differ from previous detections
                if current_detections != prev_detections:
                    for detection in current_detections:
                        # In the client script
                        try:
                            message = f"{detection[0]}, {detection[1]:.2f}"
                            client_socket.sendall(message.encode())
                        except Exception as e:
                            print(f"Failed to send message: {e}")

                    prev_detections = current_detections

            prev_frame = gray_frame

            cv2.imshow('frame', frame)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    except KeyboardInterrupt:
        print("Client disconnecting...")
        shutdown_flag.set()
    finally:
        vid.release()
        cv2.destroyAllWindows()
        client_socket.close()
        receive_thread.join()
        print("Client stopped gracefully")

if __name__ == "__main__":
    main()