import cv2
import numpy as np
from ultralytics import YOLO
import paho.mqtt.client as mqtt
from filterpy.kalman import UnscentedKalmanFilter as UKF
from filterpy.kalman import MerweScaledSigmaPoints
import time

############################################################################################################
# Initialize Variables

cy1 = 550
offset = 90 
ids = set()
pecanCount = 0 
avgPecanWeight = 0.0270453125 # lbs
refThroughput = 0 # lbs / 15 seconds
samplePeriod = 0.25 # seconds
width = 1280
height = 720
totalPecanCount = 0

# Load the YOLO11 model and set device
model = YOLO("yolo11m-pecan.pt")
device = 'cuda'

# Open the video file (use video file or webcam, here using webcam)
cap = cv2.VideoCapture('rtsp://192.168.1.10:8554/stream')

############################################################################################################
# Unscented Kalman Filter

# State transition function (identity - no control input)
def fx(x, dt):
    """ State transition function (no external control) """
    return x  # No change in state without an input

# Measurement function (identity)
def hx(x):
    """ Measurement function (direct observation) """
    return x  # We measure the state directly

points = MerweScaledSigmaPoints(n=1, alpha=0.1, beta=2, kappa=0)# Define sigma points

ukf = UKF(dim_x=1, dim_z=1, fx=fx, hx=hx, points=points, dt=samplePeriod) # Initial State Estimate
ukf.x = np.array([refThroughput / 15 / avgPecanWeight * samplePeriod]) # Initial state estimate
ukf.Q = np.array([[0.02]]) # Process noise covariance (Q) - controls how much the state changes naturally
ukf.R = np.array([[1]]) # Measurement noise covariance (R) - how noisy the measurements are
ukf.P = np.eye(1) * 0.1 # Initial state covariance (P) - initial uncertainty

############################################################################################################
# MQTT Configuration

MQTT_BROKER = "192.168.1.110"
MQTT_PORT = 1883
MQTT_TOPIC = "/jc/feedrate/"
MQTT_COUNT_TOPIC = "/jc/feedrate/count/"

def on_connect(client, userdata, flags, rc, properties=None):
    print("Connected with result code " + str(rc))
    client.subscribe(MQTT_TOPIC)

def on_message(client, userdata, message):
    global refThroughput # lbs of pecans / 15 seconds
    prevRefThroughput = refThroughput
    refThroughput = float(message.payload.decode())
    if refThroughput != prevRefThroughput:
        ukf.x = np.array([refThroughput])

client = mqtt.Client()
client.on_connect = on_connect
client.on_message = on_message

client.connect(MQTT_BROKER, MQTT_PORT, 60)
client.loop_start()  # Starts MQTT loop in the background
############################################################################################################
# Exponential Moving Average

# class FastEMA: # Exponential Moving Average Function
#     def __init__(self, alpha):
#         self.alpha = alpha
#         self.ema = None  # Initialize without a value
#         self.input = None

#     def update(self, measure, input):
#         if self.ema is None or self.input != input:
#             self.ema = measure  # Set first value
#             self.input = input
#         else:
#             self.ema = self.alpha * measure + (1 - self.alpha) * self.ema
#         return self.ema

# # Initialize the EMA filter
# ema_filter = FastEMA(alpha=0.1)  # Adjust alpha (lower = smoother)
############################################################################################################

sampleStart = time.time()
frameTime = time.time()
frameCount = 0

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

    # Define the black box position (adjust as needed)
    top_left = (0, 0) 
    bottom_right = (1280, 220)
    
    # Draw a black rectangle (filled)
    cv2.rectangle(frame, top_left, bottom_right, (0, 0, 0), thickness=-1)

    frameCount += 1 

    # sampleStart = time.time() # Sample Start Time

    # pecanCount = 0 # Reset count for new sample period

    # Run YOLO11 tracking on the frame, persisting tracks between frames
    results = model.track(frame, persist=True, classes=0, device = device)

    # Check if there are any boxes in the results
    if results[0].boxes is not None and results[0].boxes.id is not None:
        # Get the boxes (x, y, w, h), class IDs, track IDs, and confidences
        boxes = results[0].boxes.xyxy.int().cpu().tolist()  # Bounding boxes
        class_ids = results[0].boxes.cls.int().cpu().tolist()  # Class IDs
        track_ids = results[0].boxes.id.int().cpu().tolist()  # Track IDs
        confidences = results[0].boxes.conf.cpu().tolist()  # Confidence score
       
        for box, class_id, track_id, conf in zip(boxes, class_ids, track_ids, confidences):
            x1, y1, x2, y2 = box
            cy = int(y1+y2)//2
            if cy<(cy1+offset) and cy>(cy1-offset) and track_id not in ids:
                pecanCount += 1
                totalPecanCount += 1
                print(f'New Count: {pecanCount}')
                ids.add(track_id)

    # filtered_count = ema_filter.update(pecanCount, refThroughput) # Applies exponential moving average filter
    
    sampleEnd = time.time()
    print(f'Pecan Count: {pecanCount}')
    print(f'Total Count: {totalPecanCount}')
    if (sampleEnd - sampleStart) > samplePeriod:
        ukf.predict()
        print(f'Predicted State: {ukf.x[0]}')
        ukf.update(np.array([pecanCount]))
        print(f'Updated State: {ukf.x[0]}')
        filtered_count = ukf.x[0]
        print(filtered_count)
        measuredThroughput = (filtered_count * avgPecanWeight) / (samplePeriod) * 15 # lbs / 15 seconds
        print(f'Published Throughput: {measuredThroughput}')
        client.publish(MQTT_COUNT_TOPIC, str(measuredThroughput))
        pecanCount = 0
        sampleStart = time.time()
        print(f'Publish Time: {sampleStart-sampleEnd}')

    # sampleEnd = time.time() # End Sample Timer

    # samplePeriod = sampleEnd - sampleStart

    # measuredThroughput = (filtered_count * avgPecanWeight) / (samplePeriod) * 15 # lbs / 15 seconds
    
    #if (new_time := time.time()) > last_time + 0.25:
        #client.publish(MQTT_COUNT_TOPIC, str(measuredThroughput))
        #print(samplePeriod) 
        #print(str(last_time) + " " + str(new_time))
        #last_time = new_time

    if (time.time()-frameTime) > 10:
        fps = frameCount / (time.time() - frameTime)
        print(fps)
        break

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