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


# ---------------- CAMERA SETUP ----------------
picam2 = Picamera2()
picam2.set_controls({"AfMode": controls.AfModeEnum.Continuous})
picam2.start()
time.sleep(1)


# ---------------- COLOR RANGES ----------------
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])


# ---------------- 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


# ---------------- GREEN DELAY TIMER ----------------
start_time = time.time()
green_delay = 3   # seconds
following_green = False

# ---------------- 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)))


# ---------------- MAIN LOOP ----------------
try:
    while True:
        image = picam2.capture_array("main")
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        crop_img = image[250:450, 80:420]
        blur = cv2.GaussianBlur(crop_img, (5, 5), 0)
        hsv = cv2.cvtColor(blur, cv2.COLOR_BGR2HSV)

        mask_yellow = cv2.inRange(hsv, lower_yellow, upper_yellow)
        mask_green  = cv2.inRange(hsv, lower_green, upper_green)

        # ---- PERMANENT SWITCH TO GREEN ----
        elapsed_time = time.time() - start_time

        if following_green:
            mask = mask_green
            print("Following GREEN")

        elif elapsed_time > green_delay and cv2.countNonZero(mask_green) > 500:
            following_green = True
            mask = mask_green
            print("Switching to GREEN")

        else:
            mask = mask_yellow
            print("Following YELLOW")



        contours, _ = 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

                integral += error
                integral = max(min(integral, 1000), -1000)
                derivative = error - last_error
                correction = (Kp * error) + (Ki * integral) + (Kd * derivative)

                turn_slowdown = min(abs(error) * 0.1, 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

                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:
            correction = Kp * last_error
            set_motors(base_speed + correction, base_speed - correction)

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

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

        time.sleep(0.01)

except KeyboardInterrupt:
    print("\nAll done")

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