import numpy as np
import cv2
from picamera2 import Picamera2
from libcamera import controls
import time
import RPi.GPIO as GPIO


# Initialize Pi Camera
picam2 = Picamera2()
picam2.set_controls({"AfMode": controls.AfModeEnum.Continuous})
picam2.start()
time.sleep(1)  # Give camera time to start up

# Define color range for line detection (adjust these values for your specific color)
lower_color = np.array([0, 0, 0])  # Lower HSV threshold
upper_color = np.array([255, 255, 255])  # Upper HSV threshold

lower_yellow = np.array([20, 100, 100])
upper_yellow = np.array([35, 255, 255])

lower_green = np.array([40, 70, 70])
upper_green = np.array([55, 255, 255])

lower_black = np.array([90, 10, 10])
upper_black = np.array([125, 45, 45])

lower = lower_yellow
upper = upper_yellow


# ---------------- PIN CONFIG ----------------
ena, in1, in2 = 24, 26, 32
enb, in3, in4 = 31, 29, 23

GPIO.setmode(GPIO.BOARD)
GPIO.setup([ena, in1, in2, enb, in3, in4], GPIO.OUT)

motorL = GPIO.PWM(ena, 100)
motorR = GPIO.PWM(enb, 100)
motorL.start(0)
motorR.start(0)

# ---------------- PID CONSTANTS ----------------
Kp = 0.12
Ki = 0.0005
Kd = 0.06

base_speed = 25
last_error = 0
integral = 0

# ---------------- MOTOR FUNCTION ----------------
def set_motors(speed_l, speed_r):

    GPIO.output(in1, GPIO.HIGH if speed_l >= 0 else GPIO.LOW)
    GPIO.output(in2, GPIO.LOW if speed_l >= 0 else GPIO.HIGH)

    GPIO.output(in3, GPIO.HIGH if speed_r >= 0 else GPIO.LOW)
    GPIO.output(in4, GPIO.LOW if speed_r >= 0 else GPIO.HIGH)

    motorL.ChangeDutyCycle(max(0, min(abs(speed_l), 100)))
    motorR.ChangeDutyCycle(max(0, min(abs(speed_r), 100)))

try:
    while True:
        # Capture image from camera
        image = picam2.capture_array("main")
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        cv2.imshow('img', image)
        
        # Crop the image
        crop_img = image[250:450, 80:420]
        #crop_img = image[250:350, 80:180]
	
        # Apply Gaussian blur - comment out option 1 or option 2
        # Option 1: blur cropped image- use to calibrate
        #blur = cv2.GaussianBlur(crop_img, (5, 5), 0)
        
        # Option 2: blur full image- use after calibration
        blur = cv2.GaussianBlur(crop_img, (5, 5), 0)

        # Convert to HSV color space
        hsv = cv2.cvtColor(blur, cv2.COLOR_BGR2HSV)

        # Calibrate: Show HSV values
        hsv_mean = np.mean(hsv, axis=(0, 1))  # Calculate the mean HSV
        print(f"Mean HSV values: {hsv_mean}")  # Print the mean HSV values

        # Threshold the image to keep only the selected color
        mask = cv2.inRange(hsv, lower, upper)


        
        contours, hierarchy = cv2.findContours(mask.copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)

        if len(contours) > 0:

            c = max(contours, key=cv2.contourArea)

            if cv2.contourArea(c) > 800:

                M = cv2.moments(c)
                if M["m00"] != 0:
                    cx = int(M["m10"] / M["m00"])
                else:
                    continue

                frame_center = crop_img.shape[1] // 2
                error = cx - frame_center

                # -------- PID --------
                integral += error
                derivative = error - last_error

                correction = (Kp * error) + (Ki * integral) + (Kd * derivative)

                # Slow down more during sharp turns
                turn_slowdown = min(abs(error) * 0.05, 10)
                dynamic_base = base_speed - turn_slowdown

                speed_l = dynamic_base + correction
                speed_r = dynamic_base - correction

                set_motors(speed_l, speed_r)

                last_error = error

                # Debug lines
                cv2.line(crop_img, (cx, 0), (cx, crop_img.shape[0]), (255, 0, 0), 2)
                cv2.line(crop_img, (frame_center, 0), (frame_center, crop_img.shape[0]), (0, 255, 0), 2)
                cv2.drawContours(crop_img, [c], -1, (0, 255, 0), 2)

        else:
            # If line lost, keep turning based on last direction
            correction = Kp * last_error
            set_motors(base_speed + correction, base_speed - correction)

        cv2.imshow("Camera", crop_img)
        cv2.imshow("Mask", mask)

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

        time.sleep(0.01)



        # Show HSV and Mask images for debugging
        cv2.imshow('hsv', hsv)  # Show the HSV image
        cv2.imshow('mask', mask)  # Show the thresholded mask

        cv2.waitKey(1)
        
except KeyboardInterrupt:
    print('\nAll done')
    pass


finally:
    motorL.stop()
    motorR.stop()
    GPIO.cleanup()
    cv2.destroyAllWindows()