import RPi.GPIO as GPIO
import time

# Setup GPIO and pins for color sensor
s0 = 13
s1 = 15
s2 = 16
s3 = 18
sig = 22
cycles = 10
GPIO.setmode(GPIO.BOARD)
GPIO.setup(s0, GPIO.OUT)
GPIO.setup(s1, GPIO.OUT)
GPIO.setup(s2, GPIO.OUT)
GPIO.setup(s3, GPIO.OUT)
GPIO.setup(sig, GPIO.IN, pull_up_down=GPIO.PUD_DOWN)

# Set frequency scaling
GPIO.output(s0, GPIO.HIGH)
GPIO.output(s1, GPIO.LOW)

# color sensor LED
GPIO.setup(37, GPIO.OUT)
GPIO.output(37, GPIO.HIGH)

# stepper motor setup
# Define the GPIO pins for the L298N motor driver
OUT1 = 32
OUT2 = 29
OUT3 = 31
OUT4 = 33

# Set the GPIO pins as output
GPIO.setup(OUT1, GPIO.OUT)
GPIO.setup(OUT2, GPIO.OUT)
GPIO.setup(OUT3, GPIO.OUT)
GPIO.setup(OUT4, GPIO.OUT)

# Set initial state of pins to low
GPIO.output(OUT1,GPIO.LOW)
GPIO.output(OUT2,GPIO.LOW)
GPIO.output(OUT3,GPIO.LOW)
GPIO.output(OUT4,GPIO.LOW)

# Sets number of steps to move
num_steps = 25

# Sets delay between steps
step_delay = 0.03

# set up servo motor
GPIO.setup(11,GPIO.OUT)
servo = GPIO.PWM(11, 50)
servo.start(0)  
servo_position = "r" 

# average what the color sensor senses for rgb over a set number of cycles
def DetectColor():
    # Detect red values
    GPIO.output(s2, GPIO.LOW)
    GPIO.output(s3, GPIO.LOW)
    time.sleep(0.2)
    start_time = time.time()
    for count in range(cycles):
        GPIO.wait_for_edge(sig, GPIO.FALLING)
    duration = time.time() - start_time
    red = cycles / duration
    
    # Detect blue values
    GPIO.output(s2, GPIO.LOW)
    GPIO.output(s3, GPIO.HIGH)
    time.sleep(0.1)
    start_time = time.time()
    for count in range(cycles):
        GPIO.wait_for_edge(sig, GPIO.FALLING)
    duration = time.time() - start_time
    blue = cycles / duration
    
    # Detect green values
    GPIO.output(s2, GPIO.HIGH)
    GPIO.output(s3, GPIO.HIGH)
    time.sleep(0.2)
    start_time = time.time()
    for count in range(cycles):
        GPIO.wait_for_edge(sig, GPIO.FALLING)
    duration = time.time() - start_time
    green = cycles / duration
    rgb = [red, green, blue]
    return rgb


# check which recorded color a new color is closest to
def determineColor(reds, greens, blues, rgb):
    red_diffs = []
    yellow_diffs = []
    green_diffs = []
    blue_diffs = []
    for i in range(3):
        red_diffs.append(abs(rgb[i]-reds[i]))
        yellow_diffs.append(abs(rgb[i]-yellows[i]))
        green_diffs.append(abs(rgb[i]-greens[i]))
        blue_diffs.append(abs(rgb[i]-blues[i]))
    print("a", rgb)
    print("r", red_diffs)
    print("y", yellow_diffs)
    print("g", green_diffs)
    print("b", blue_diffs)
    red_sum = red_diffs[0] + red_diffs[1] + red_diffs[2]
    yellow_sum = yellow_diffs[0] + yellow_diffs[1] + yellow_diffs[2]
    green_sum = green_diffs[0] + green_diffs[1] + green_diffs[2]
    blue_sum = blue_diffs[0] + blue_diffs[1] + blue_diffs[2]
    data = sorted([red_sum, yellow_sum, green_sum, blue_sum])
    if data[0] == red_sum:
        print("red")
        return "r"
    elif data[0] == yellow_sum:
        print("yellow")
        return "y"
    elif data[0] == green_sum:
        print("green")
        return "g"
    elif data[0] == blue_sum:
        print("blue")
        return "b"
    else:
        print("what?")
    return


# rotate the stepper motor clockwise a given number of steps
def moveStepper(num_steps):
    current_step = 0
    step_delay = 0.03
    for x in range(num_steps):
        if current_step == 0:
            GPIO.output(OUT1,GPIO.HIGH)
            GPIO.output(OUT2,GPIO.LOW)
            GPIO.output(OUT3,GPIO.HIGH)
            GPIO.output(OUT4,GPIO.LOW)
            time.sleep(step_delay)
            #print("step 0")
        elif current_step == 1:
            GPIO.output(OUT1,GPIO.HIGH)
            GPIO.output(OUT2,GPIO.LOW)
            GPIO.output(OUT3,GPIO.LOW)
            GPIO.output(OUT4,GPIO.HIGH)
            time.sleep(step_delay)
            #print("step 1")
        elif current_step == 2:
            GPIO.output(OUT1,GPIO.LOW)
            GPIO.output(OUT2,GPIO.HIGH)
            GPIO.output(OUT3,GPIO.LOW)
            GPIO.output(OUT4,GPIO.HIGH)
            time.sleep(step_delay)
        elif current_step == 3:
            GPIO.output(OUT1,GPIO.LOW)
            GPIO.output(OUT2,GPIO.HIGH)
            GPIO.output(OUT3,GPIO.HIGH)
            GPIO.output(OUT4,GPIO.LOW)
            time.sleep(step_delay)
        if current_step == 3:
            current_step = 0
            continue 
        current_step = current_step + 1
    return


# move the servo motor to the position related to each color's basket
def moveServo(color):
    if color == "r":
        servo.ChangeDutyCycle(2.5)
    elif color == "y":
        servo.ChangeDutyCycle(5)
    elif color == "g":
        servo.ChangeDutyCycle(8)
    elif color == "b":
        servo.ChangeDutyCycle(12.5)
    else:
        print("ERROR: SERVO DID NOT RECIEVE VALID COLOR")


# main code, read in initial colors, enter a loop to read new colors and drop the balls
# in the correct baskets
try:
    initialized = False
    reds = []
    yellows = []
    greens = []
    blues = []
    while not initialized:
        # get red
        redsamples = []
        for i in range(3):
            redsamples.append(DetectColor())
            time.sleep(.01)
        for i in range(3):
            reds.append((redsamples[0][i] + redsamples[1][i] + redsamples[2][i])/3)
        print("reds:", reds)
        moveStepper(25)

        # get green
        greensamples = []
        for i in range(3):
            greensamples.append(DetectColor())
            time.sleep(.01)
        for i in range(3):
            greens.append((greensamples[0][i] + greensamples[1][i] + greensamples[2][i])/3)
        print("greens:", greens)
        moveServo("r")
        moveStepper(25)

        # get blue
        bluesamples = []
        for i in range(3):
            bluesamples.append(DetectColor())
            time.sleep(.01)
        for i in range(3):
            blues.append((bluesamples[0][i] + bluesamples[1][i] + bluesamples[2][i])/3)
        print("blues:", blues)
        moveServo("g")
        moveStepper(25)

        # get yellow
        yellowsamples = []
        for i in range(3):
            yellowsamples.append(DetectColor())
            time.sleep(.01)
        for i in range(3):
            yellows.append((yellowsamples[0][i] + yellowsamples[1][i] + yellowsamples[2][i])/3)
        print("yellows:", yellows)
        moveServo("b")
        moveStepper(25)
        print("initialized")
        initialized = True
        
    color0 = "y"
    while True:
        print("detecting current color in 3")
        time.sleep(1)
        print("2")
        time.sleep(1)
        print("1")
        time.sleep(1)
        rgb = DetectColor()
        color1 = determineColor(reds, greens, blues, rgb)
        print("moving servo")
        moveServo(color0)
        color0 = color1
        print("moving stepper")
        moveStepper(25)


except KeyboardInterrupt:
    GPIO.cleanup()