import cv2
import numpy as np
import serial
import time

# =====================
# CONFIGURATION
# =====================
CAMERA_INDEX = 0
SERIAL_PORT = "COM3"  # Change to your Arduino port
BAUD_RATE = 9600

# =====================
# ARDUINO SETUP
# =====================
try:
    arduino = serial.Serial(SERIAL_PORT, BAUD_RATE, timeout=1)
    time.sleep(2)
    print("✅ Connected to Arduino!")
except:
    arduino = None
    print("⚠️  Running in demo mode")

def send_command(cmd):
    """Send command to Arduino"""
    if arduino and arduino.is_open:
        arduino.write((cmd + "\n").encode("utf-8"))
        print(f"✅ Sent: {cmd}")
    else:
        print(f"⚠️  Would send: {cmd}")

# =====================
# VISION FUNCTIONS
# =====================
def nothing(x):
    pass

def get_mask(frame, hsv_ranges):
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    lower = np.array([hsv_ranges["H Min"], hsv_ranges["S Min"], hsv_ranges["V Min"]])
    upper = np.array([hsv_ranges["H Max"], hsv_ranges["S Max"], hsv_ranges["V Max"]])
    mask = cv2.inRange(hsv, lower, upper)
    return mask

def get_contour_info(mask, frame):
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return frame, None

    c = max(contours, key=cv2.contourArea)
    if cv2.contourArea(c) < 500:
        return frame, None

    cv2.drawContours(frame, [c], -1, (0, 0, 255), 3)

    x_vals = c[:, 0, 0]
    y_vals = c[:, 0, 1]
    w_px = max(x_vals) - min(x_vals)
    h_px = max(y_vals) - min(y_vals)

    if w_px == 0 or h_px == 0:
        return frame, None

    oar = round(float(w_px) / h_px, 2)
    return frame, oar

# =====================
# MAIN PROGRAM
# =====================
cap = cv2.VideoCapture(CAMERA_INDEX)

# Create windows
cv2.namedWindow("Trackbars", cv2.WINDOW_NORMAL)
cv2.namedWindow("Original", cv2.WINDOW_NORMAL) 
cv2.namedWindow("Mask", cv2.WINDOW_NORMAL)
cv2.namedWindow("Result", cv2.WINDOW_NORMAL)

# Create HSV trackbars
for name, maxval, initval in [
    ("H Min", 179, 0), ("S Min", 255, 0), ("V Min", 255, 0),
    ("H Max", 179, 179), ("S Max", 255, 255), ("V Max", 255, 255)
]:
    cv2.createTrackbar(name, "Trackbars", initval, maxval, nothing)

last_shape = None

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

    # Get HSV values from trackbars
    hsv_ranges = {name: cv2.getTrackbarPos(name, "Trackbars") for name in
                  ["H Min", "S Min", "V Min", "H Max", "S Max", "V Max"]}

    # Process frame
    mask = get_mask(frame, hsv_ranges)
    result, oar = get_contour_info(mask, frame.copy())

    # =====================
    # SIMPLE IF-ELSE - EDIT THIS PART!
    # =====================
    if oar is None:
        shape = "No object"
        current_cmd = None
        
    # EDIT THESE CONDITIONS BELOW:
    elif oar >= 1.5:
        shape = "Very Square"
        if last_shape != shape:
            send_command("S1180")
            time.sleep(0.1)
            send_command("S2180")
            last_shape = shape
            
    elif oar >= 1.2:
        shape = "Square" 
        if last_shape != shape:
            send_command("S1135")
            time.sleep(0.1)
            send_command("S2135")
            last_shape = shape
            
    elif oar >= 0.8:
        shape = "Rectangle"
        if last_shape != shape:
            send_command("S190")
            time.sleep(0.1)
            send_command("S290")
            last_shape = shape
            
    elif oar >= 0.5:
        shape = "Long Rectangle"
        if last_shape != shape:
            send_command("S145")
            time.sleep(0.1)
            send_command("S245")
            last_shape = shape
            
    else:
        shape = "Very Long Rectangle"
        if last_shape != shape:
            send_command("S10")
            time.sleep(0.1)
            send_command("S20")
            last_shape = shape

    # Display info on screen
    if oar:
        cv2.putText(result, f"OAR={oar}", (30, 50),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    cv2.putText(result, shape, (30, 90),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)

    # Show all windows
    cv2.imshow("Original", frame)
    cv2.imshow("Mask", mask)
    cv2.imshow("Result", result)

    # Press 'q' to quit
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# Clean up
cap.release()
cv2.destroyAllWindows()
if arduino:
    arduino.close()