# 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
# Broadcast on channel 3
hub = PrimeHub(broadcast_channel=3)

# Initialize robot constants

# Initialize motors
gate_left = Motor(Port.A, Direction.COUNTERCLOCKWISE)
gate_right = Motor(Port.B, Direction.CLOCKWISE)
lights_right = Motor(Port.C, Direction.CLOCKWISE) # Gate lights
lights_left = Motor(Port.D, Direction.COUNTERCLOCKWISE) # Volcano lights

# Initialize sensors
touch_sensor = ForceSensor(Port.E)
us_sensor = UltrasonicSensor(Port.F)

# Initialize program constants
# Angle to open gate, to turn light axle, motor speed in deg/s, time to wait in ms, list of motors
angle_gate = -90
angle_lights = 360
speed_motor = 720
wait_time = 100
list_motors = [gate_left,gate_right,lights_right,lights_left,'',''] # Port A-F

# Declare program variables
# Initial frequency of lights flickering, initial message, initial US sensor counter, flag for stopping
frequency = 'low'
message = 'start'
counter_us = 0
not_stopped = True

# Declare functions
# Start lights at frequency of change
def light_up(frequency):

    #while True:
    print('Turning lights on at', frequency, 'frequency...')
    # If frequency low, choose high wait time in ms
    if frequency == 'low':
        wait_lights = randint(200,500)
    # If frequency high, choose low wait time in ms, run volcano lights
    else:
        wait_lights = randint(100,200)
        #lights_left.run_angle(speed_motor, angle_lights, wait=False) # Volcano lights
    # Run gate lights
    lights_right.run_angle(speed_motor, angle_lights) # Gate lights
    wait(wait_lights)
    lights_right.hold() # Gate lights
    #lights_left.hold() # Volcano lights
    wait(wait_time)

# Open or close gate
def open_close_gate(action):
    # Set as global to allow changing
    global angle_gate
    speed = speed_motor / 5
    if action == 'open':
    	print('Opening gate...')
    	gate_left.run_angle(speed, angle_gate, wait=False)
    	gate_right.run_angle(speed, angle_gate)
        angle_gate = angle_gate * -1
    if action == 'close':
    	print('Closing gate...')
    	gate_left.run_angle(speed, angle_gate, wait=False)
        gate_right.run_angle(speed, angle_gate)
        angle_gate *= -1
    else:
	    pass

# Reset gate & light 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

# Send message
def send_message(message):
    print('Sending message', message,'...')
    # Change button colour, send message
    hub.light.on(Color.RED)
    hub.ble.broadcast(message)
    # Change button colour
    hub.light.on(Color.GREEN)
    wait(wait_time)

# Declare main program
# Reset gate & lights motors, set angle for all active motors
motors_start = [0,0,270,0,'',''] # Port A-F
motors_stop = [0,0,0,0,'',''] # Port A-F
reset(motors_start)

while not_stopped:

    # Start lights at low frequency of change
    light_up(frequency)
    # If object detected less than 275mm or touch sensor pressed
    if us_sensor.distance() < 275 or touch_sensor.touched():
        counter_us += 1
        temp_time = wait_time * 10
        print('US counter', counter_us)
        wait(temp_time)
        # Detect Maddie on US sensor, open gate, send message to start Dig & Nest robots
        if counter_us == 1:
            send_message(message)
            # Send message to prepare Dig robot for waving
            message = 'wave'
            action = 'open'
            send_message(message)
            # Open gate
            open_close_gate(action)
        # Detect Jeep on US sensor
        elif counter_us == 2:
            wait(temp_time)
            action = 'close'
            #send_message(message)
            # Close gate
            open_close_gate(action)
        # Detect Em on touch sensor, change message, send message to start Dig & Nest robots shaking a little
        elif counter_us == 3:
            message = 'rumble'
            #action = 'open'
            send_message(message)
            # Open gate
            #open_close_gate(action)
            light_up(frequency)
        # Detect T-Rex on US sensor, change message, send message to start Dig & Nest robots shaking a lot
        elif counter_us == 4:
            message = 'trex_attack'
            #action = 'close'
            frequency = 'high'
            send_message(message)
            wait(temp_time)
            # Close gate
            #open_close_gate(action)
            # Start lights at high frequency of change
            light_up(frequency)
        # Detect Em on US sensor, change message, send message to stop Dig & Nest robots
        elif counter_us == 5:
            message = 'stop'
            send_message(message)
            # Reset gate & lights motors, set angle for all active motors
            reset(motors_stop)
            #wait(temp_time)
            not_stopped = False
        else:
            pass