# Import modules
from pybricks.hubs import PrimeHub
from pybricks.parameters import Button, Color, Direction, Port, Side, Stop
from pybricks.pupdevices import ColorSensor, ForceSensor, Motor, UltrasonicSensor
from pybricks.robotics import DriveBase
from pybricks.tools import multitask, run_task, StopWatch, wait
from urandom import randint

# Initialize brick
# Listen on channel 1
hub = PrimeHub(observe_channels=[3])

# Initialize robot constants

# Initialize motors
moses_bottom = Motor(Port.A, Direction.CLOCKWISE) # Moses' head
pedro_bottom = Motor(Port.B, Direction.CLOCKWISE) # Pedro's lift
moses_top = Motor(Port.C, Direction.CLOCKWISE) # Moses' mouth
pedro_top = Motor(Port.D, Direction.CLOCKWISE) # Pedro's wings
#tree_right = Motor(Port.F, Direction.CLOCKWISE) # Right tree motor

# Initialize sensors
us_sensor = UltrasonicSensor(Port.E)

# Initialize program constants
# Open, lift, rotate & flap angle, motor speed in deg/s, time to wait in ms, list of motors
angle_open = 120
angle_lift = 120
angle_rotate = 60
#angle_shake = 30
angle_flap = 360
speed_motor = 360
wait_time = 100
list_motors = [moses_bottom,pedro_bottom,moses_top,pedro_top,'',''] # Port A-F

# Declare program variables
# Initial frequency of flapping, blank message, flag for lifting & stopping
frequency = 'low'
message = ''
not_lifted = True
not_stopped = True

# Declare functions
# Receive message
def receive_message():
    # Initialize flag
    not_received = True

    while not_received:
        # Listen for message, turn light off
        message = hub.ble.observe(3)
	# Test code parts with hardcoded message
        #message = 'start'
        #message = 'wave'
        #message = 'rumble'
        #message = 'trex_attack'
        #message = 'stop'
        hub.light.off()
        # If no message received, turn light red
        if message is None:
            hub.light.on(Color.RED)
        # If message received, turn light green, change flag
        else:
            print('Receiving message', message, '...')
            hub.light.on(Color.GREEN)
            not_received = False
	    # Wait until next message check
        wait(wait_time)

    return message

# Reset Moses, Pedro & Tree motors
def reset(motors):
    # Initialize counter
    counter_motor = 0
    # Loop through motors array
    for angle in motors:
	    # If a value is found in array
        if angle != '':
            print('Resetting motor', list_motors[counter_motor], 'at', angle, 'degrees...')
            # Look up motor, reset position to angle
            list_motors[counter_motor].run_target(speed_motor, angle)
        # Increment counter
	    counter_motor += 1

# Flap wings & rotate head at frequency of change
def flap_rotate(frequency):
    # Initialize angles, counter, range
    temp_rotate = angle_rotate * -1 / 2
    temp_open = angle_open * -1 / 2
    temp_flap = angle_flap * -1 / 2
    counter = 0
    range = 2
    print('Flapping & rotating at', frequency, 'frequency ...')
    # Reset wings & head motors at same time
    moses_bottom.run_target(speed_motor, temp_rotate, wait=False) # Rotate motor
    moses_top.run_target(speed_motor, temp_open, wait=False) # Open motor
    pedro_top.run_target(speed_motor, temp_flap) #, wait=False) # Flap motor
    temp_rotate = angle_rotate
    temp_open = angle_open
    temp_flap = angle_flap
    # Repeat 2 times
    for counter in range(range):
        # If frequency low, choose low speed in deg/s
        if frequency == 'low':
            speed_flap = randint(50,100)
            speed_rotate = randint(50,100)
        # If frequency high, choose high speed in deg/s 
        else:
            speed_flap = randint(100,200)
            speed_rotate = randint(100,200)
        # Rotate head & flap wings, run motors at same time
        moses_bottom.run_angle(speed_rotate, temp_rotate, wait=False) # Rotate motor
        moses_top.run_target(speed_rotate, temp_open, wait=False) # Open motor
        pedro_top.run_angle(speed_flap, temp_flap) #, wait=False) # Flap motor
        # Toggle angle to move motors back & forth
        temp_rotate = temp_rotate * -1
        temp_open = temp_open * -1
        temp_flap = temp_flap * -1
        counter += 1     
       
# Make Pedro flap & lift
def flap_lift(action):
    # Set as global to allow changing
    global angle_flap, angle_lift
    # Initialise high speed, counter, range
    speed = speed_motor
    counter = 0
    range = 2
    print('Making Pedro', action, '...')
    # Reset Pedro motors
    if action == 'flap':
        range = 6
        pedro_top.run_target(speed, angle_flap) # Flap motor
    elif action == 'lift':
        pedro_bottom.run_target(speed, angle_lift) # Lift motor
    elif action == 'flap_lift':
        pedro_top.run_target(speed, angle_flap, wait=False) # Flap motor
        pedro_bottom.run_target(speed, angle_lift) # Lift motor
    else:
        pass
    # Repeat 2 or 6 times
    for counter in range(range):
        # Lift 1st
        #angle_flap = angle_flap * -1
        if action == 'flap':
            pedro_top.run_angle(speed, angle_flap) # Flap motor
        elif action == 'lift':
            pedro_bottom.run_angle(speed, angle_lift) # Lift motor
        elif action == 'flap_lift':
            pedro_top.run_angle(speed, angle_flap, wait=False) # Flap motor
            pedro_bottom.run_angle(speed, angle_lift) # Lift motor
        else:
            pass
        counter += 1
    # Toggle angle to move motor back & forth
    angle_flap = angle_flap * -1
    angle_lift = angle_lift * -1
    # Reset Pedro motors
    if action == 'flap':
        pedro_top.run_target(speed, angle_flap) # Flap motor
    elif action == 'lift':
        pedro_bottom.run_target(speed, angle_lift) # Lift motor
    elif action == 'flap_lift':
        pedro_top.run_target(speed, angle_flap, wait=False) # Flap motor
        pedro_bottom.run_target(speed, angle_lift) #, wait=False) # Lift motor
    else:
        pass

# Declare main program
# Reset Moses, Pedro motors, set angle for all active motors
motors = [0,-90,-90,0,'',''] # Port A-F
reset(motors)

while not_stopped:

    # Listen to message
    message = receive_message()
    # If message 'start', make Pedro flap
    if message == 'start':
        action = 'flap'
        flap_lift(action)
    # If message 'wave', if object detected less than 200mm
    elif message == 'wave':
        # Detect Jeep on US sensor, make Pedro flap & lift but only 1 time
        if us_sensor.distance() < 200:
            if not_lifted:
                action = 'flap_lift'
                flap_lift(action)
                not_lifted = False
                #action = 'flap'
                #flap_lift(action)
        # If no object detected less than 200mm continue flapping until lifted
        #else:
        #    if not_lifted:
        #        action = 'flap'
        #        flap_lift(action)
    # If message 'rumble', flap wings & rotate head a little
    elif message == 'rumble':
        flap_rotate(frequency)
    # If message 'T-Rex attack', flap wings & rotate head a lot, make Pedro flap & lift
    elif message == 'trex_attack':
        frequency = 'high'
        action = 'flap_lift'
        # Run both functions at same time
        #multitask(
        #    run_task(flap_rotate(frequency)),
        #    run_task(flap_lift(action))
        #)
        flap_rotate(frequency)
        flap_lift(action)
    # If message 'stop', stop Moses, Pedro, reset motors
    elif message == 'stop':
        reset(motors)
        not_stopped = False    
    else:
        pass