#!/usr/bin/env pybricks-micropython

# Build loosely based on RobotEducator robot code and building instructions:
# https://github.com/pybricks/pybricks-projects/blob/master/sets/mindstorms-ev3/education-core/robot_educator_line/main.py
# https://github.com/pybricks/pybricks-projects/blob/master/sets/mindstorms-ev3/education-core/robot_educator_ultrasonic/main.py
# https://assets.education.lego.com/v3/assets/blt293eea581807678a/blt8b300493e30608e9/5f8801dfb8b59a77a945d13c/ev3-rem-color-sensor-down-driving-base.pdf?locale=en-us
# https://assets.education.lego.com/v3/assets/blt293eea581807678a/blte04fb7bf4f716f3d/5f8801e3bf5ab07ee90076c9/ev3-ultrasonic-sensor-driving-base.pdf?locale=en-us

# Main changes include:
# - Combining ultrasonic sensor and 2 colour sensors on chassis
# - Adding touch sensor for colour sensor calibration
# - Modifying code to use ultrasonic sensor and 2 colour sensors
# - Adding Wi-Fi 2-way connectivity to start 3 other robots and play their sounds using Jeep's Bluetooth mic
# - Using car sounds and display image (in .png format) created from stock
# - Adding 'La Cucaracha' sound from https://freesound.org/people/csnmedia/sounds/388276/ .

# Import modules
import math, socket, sys, _thread, time, urandom

from pybricks.hubs import EV3Brick

from pybricks.ev3devices import ColorSensor, GyroSensor, InfraredSensor, Motor, TouchSensor, UltrasonicSensor

from pybricks.media.ev3dev import Font, ImageFile, SoundFile
from pybricks.parameters import Button, Color, Direction, Port, Stop
from pybricks.robotics import DriveBase
from pybricks.tools import DataLog, StopWatch, wait 

from ucollections import namedtuple

# Initialize brick
Jeep = EV3Brick()

# Initialize robot constants
axle_track = 153 # mm distance between middle of tire contact
wheel_diameter = 55.5 # mm

font_big = Font(size=24)
font_small = Font(size=8)

Jeep.screen.set_font(font_small)

# Initialize motors
#motor_up = Motor(Port.A)
motor_left = Motor(Port.B)
motor_right = Motor(Port.C)
#motor_down = Motor(Port.D)

drive_base = DriveBase(motor_left, motor_right, wheel_diameter, axle_track)

# Initialize sensors
sensor_colour_left = ColorSensor(Port.S1)
sensor_colour_right = ColorSensor(Port.S2)
#sensor_gyro = GyroSensor(Port.S2)
#sensor_ir = InfraredSensor(Port.S4)
sensor_touch = TouchSensor(Port.S3)
sensor_us = UltrasonicSensor(Port.S4)

# Initialize program constants
# Number of required identical colour detections
DEBOUNCE_FRAMES = 4
# Colour value & name dictionary to identify habitats / scenes
COLOURS = { 
    Color.GREEN: 'green', # Brontie - Forest
    Color.YELLOW: 'yellow', # Emet - Desert
    Color.BLUE: 'blue', # Trixie - River / lake
    Color.RED: 'red'} # Rexie - Danger

COLOUR_RED = 'red' # T-Rex attack

RATE_HIGH = 75 # deg/s
RATE_LOW = 20
RATE_MEDIUM = 50

SPEED_HIGH = 100 # mm/s
SPEED_LOW = 20
SPEED_MEDIUM = 50

THRESHOLD_OBSTACLE = 200 # mm

TIME_WAIT = 1 # ms

# Declare program variables
# Flags for colours found
colour_blue_not_found = True
colour_green_not_found = True
colour_red_not_found = True
colour_yellow_not_found = True
colour_blue_1st = True
colour_red_1st = True

zone = ""
counter_red = 0

intensity_left = 0
intensity_right = 0
left_right_white = False
speed_toggle = False

# Flag for looking for zones
looking_for_zones = True
# Flag for playing sound
playing_sound = False
# Flag for T-Rex not detected
trex_not_detected = True

# Buffer to hold stable detections
colour_buffer = []
colour_last = None

# Declare functions
# Avoid obstacle by backing up
def avoid_obstacle():
    # Declare function variables
    #global trex_not_detected, playing_sound
    # Initialize function constants
    DISTANCE = 200 # mm
    # Detect T-Rex, stop to talk, play sound
    #playing_sound = True
    #print("Sound busy speaking DANGER")
    #Jeep.speaker.say("Danger, Danger, Rex See the T-Rex is attacking, Backing up, I'm backing up")
    #Jeep.speaker.play_file('Backing_Alert.wav')
    # Move backward a little, stop to talk
    drive_base.straight(-DISTANCE)
    drive_base.stop()
    #Jeep.speaker.say('Abandon vehicle, Powering down')
    #Jeep.speaker.play_file('Power_Down.wav')
    #send_message('rexie')
    #playing_sound = False

# Calibrate colour sensors
def calibrate():
    # Declare function variables
    global left_black, left_black_max, left_range, left_white, left_white_min
    global right_black, right_black_max, right_range, right_white, right_white_min
    global threshold
    # Calibrate
    Jeep.speaker.say('Calibrating black')
    # Wait until touch sensor pressed
    while not sensor_touch.pressed():
        pass

    left_black = sensor_colour_left.reflection()
    right_black = sensor_colour_right.reflection()

    Jeep.speaker.say('Calibrating white')
    # Wait until touch sensor pressed
    while not sensor_touch.pressed():
        pass

    left_white = sensor_colour_left.reflection()
    right_white = sensor_colour_right.reflection()

    # Calculate variables
    if left_white >= right_white:
        min_white = right_white
    else:
        min_white = left_white

    if left_black >= right_black:
        max_black = left_black
    else:
        max_black = right_black

    threshold = ( min_white + max_black ) / 2

    print('Left white:', left_white, 'Right white:', right_white, 'White min:', min_white)
    print('Left black:', left_black, 'Right black:', right_black, 'Black max:', max_black)
    print('Threshold :', threshold)

# Detect colours
def detect_colours():
    # Declare function variables
    global colour_blue_not_found, colour_green_not_found, colour_red_not_found, colour_yellow_not_found
    global colour_blue_1st, colour_red_1st
    global colour_last, colour_buffer, msg, zone, trex_not_detected, counter_red
    # Get colour readings & find name from dictionary
    colour_left = sensor_colour_left.color()
    colour_right = sensor_colour_right.color()

    colour_left_name = COLOURS.get(colour_left, 'Unknown')
    colour_right_name = COLOURS.get(colour_right, 'Unknown')

    # Exit if both sensors do not see same colour (for more stability)
    if colour_left != colour_right:
        colour_buffer.clear()
        return "Unknown"

    colour_current = colour_left

    # Check for repeat colour detection
    if colour_current == colour_last:
        colour_buffer.append(colour_current)
    else:
        colour_buffer.clear()
        colour_last = colour_current

    # Exit if colour not yet stable
    if len(colour_buffer) < DEBOUNCE_FRAMES:
        return "Unknown"

    # Continue since colour confirmed stable
    # If red and not found yet, increase counter, set zone
    if colour_current == Color.RED and counter_red == 0:
        counter_red = 1
        zone = "FOREST"
        #send_message('brontie')

    # If red and found 1 time, increase counter, set zone
    elif colour_current == Color.RED and counter_red == 1:
        counter_red = 2
        zone = "DESERT"
        
    # If red and found 2 times, increase counter, set zone
    elif colour_current == Color.RED and counter_red == 2:
        counter_red = 3
        zone = "RIVER"
        #send_message('trixie')

    # If red and found 3 times, set flag & zone
    elif colour_current == Color.RED and counter_red == 3:
        # Let T-Rex reaction handle this
        #colour_red_not_found = False
        trex_not_detected = False
        zone = "DANGER"

    # Reset buffer after valid detection
    colour_buffer.clear()
    return colour_current

# Drive
def drive():
    # Declare function variables
    global left_right_white, threshold
    # Get reflection readings
    intensity_left = sensor_colour_left.reflection()
    intensity_right = sensor_colour_right.reflection()

    #print('Left intensity:', intensity_left, 'Right intensity:', intensity_right, 'Threshold:', threshold)

    if intensity_left >= threshold and intensity_right >= threshold:
        # Both see white, move forward
        drive_base.drive(SPEED_MEDIUM, 0)
    elif intensity_left < threshold and intensity_right >= threshold:
        # Left sensor sees black, turn left medium
        drive_base.drive(SPEED_MEDIUM, -RATE_MEDIUM)
    elif intensity_left >= threshold and intensity_right < threshold:
        # Right sensor sees black, turn right medium
        drive_base.drive(SPEED_MEDIUM, RATE_MEDIUM)
    elif intensity_left < threshold and intensity_right < threshold:
        # Both see 'black', move forward slowly
        drive_base.drive(SPEED_LOW, 0)
    # Wait to avoid excessive loop execution
    wait(TIME_WAIT)

# Send message to dinosaur
def send_message(dinosaur):

    host = 'jeep'  # Jeep
    #host = '192.168.0.119'  # Jeep
    port = 12345

    dinosaurs = ['brontie', 'trixie', 'rexie']  # Brontie, Trixie, Rexie
    #dinosaurs = ['192.168.0.176', '192.168.0.149', '192.168.0.159']  # Brontie, Trixie, Rexie

    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    # If no response received in 2 sec, move on
    s.settimeout(2)

    try:
        s.connect((dinosaur, port))
        s.sendall(b'START')
        print("Message START sent to", dinosaur, "successfully.")
    except OSError:
        print("Connection to", dinosaur, "unsuccessful! Timeout!")
    except socket.error as e:
        print("Connection to", dinosaur, "unsuccessful! Error:", e)

    s.close()

# Speak message for zone
def speak_message():

    global zone, msg, playing_sound
    if zone != "":

        playing_sound = True
        print("Zone", zone, "entered. Sound busy speaking.")

    if zone == "START":
        Jeep.speaker.say('Welcome to Jurassic Park, Keep limbs inside vehicle at all times, Some dinosaurs are carnivore')
    elif zone == "DESERT":

        Jeep.speaker.say("Entering desert district, Home of Billy Bob, Our resident digger")
        Jeep.speaker.play_file("La_Cucaracha.wav")
        print("Playing 'La_Cucaracha.wav'")

    elif zone == "FOREST":

        Jeep.speaker.say("Entering forest district, Home of Brontie, the Bronto, Oops, The Brachiosaurus")
        send_message('brontie')

    elif zone == "RIVER":

        Jeep.speaker.say("Entering lake district, Home of Trixie, the Triceratops")
        send_message('trixie')

    elif zone == "DANGER":

        Jeep.speaker.say("Danger, Danger, Rex See the T-Rex is attacking, Backing up, I'm backing up")
        Jeep.speaker.play_file('Backing_Alert.wav')
        print("Playing 'Backing_Alert.wav'")
        send_message('rexie')

    elif zone == "END":

        Jeep.speaker.say('Abandon vehicle, Powering down')
        Jeep.speaker.play_file('Power_Down.wav')
        print("Playing 'Power_Down.wav'")
        #send_message('rexie')

    else:
        pass
    zone = ""
    playing_sound = False
    #print("Sound not busy speaking")

# Receive message from dinosaur
def receive_message():

    global msg, playing_sound
    #if msg == "DIG":

        #playing_sound = True
        #Jeep.speaker.play_file("La_Cucaracha.wav")
        #print("Playing 'La_Cucaracha.wav'")
        #msg = "BLANK"
        #playing_sound = False

    #elif playing_sound:
    if playing_sound:
        pass
    else:

        try:
            conn, addr = r.accept()
            msg = conn.recv(1024).decode()

            if msg != "":

                playing_sound = True
                print("Message", msg, "received successfully.")

            if msg == "GROWL":

                Jeep.speaker.play_file("Rexie_Growl.wav")
                print("Playing 'Rexie_Growl.wav'...")

            elif msg == "ATTACK":

                Jeep.speaker.play_file("Rexie_Attack.wav")
                print("Playing 'Rexie_Attack.wav'...")

            elif msg == "WIN":

                Jeep.speaker.play_file("Rexie_Win.wav")
                print("Playing 'Rexie_Win.wav'...")

            elif msg == "MOAN":

                Jeep.speaker.play_file("Brontie_Moan.wav")
                print("Playing 'Brontie_Moan.wav'...")

            elif msg == "HAPPY":

                Jeep.speaker.play_file("Trixie_Happy.wav")
                print("Playing 'Trixie_Happy.wav'...")

            elif msg == "ANGRY":

                Jeep.speaker.play_file("Trixie_Angry.wav")
                print("Playing 'Trixie_Angry.wav'...")

            else:
                print("Message", msg, "unrecognised!")
            msg = "BLANK"
            
            conn.close()

        except OSError:
            pass
            #print("OS error!")
        except socket.error as e:
            print("Socket error! Error:", e)
        playing_sound = False
        #print("Sound not busy receiving")

# Main program
# Calibrate colour sensors
calibrate()

# Set Jeep image
Jeep.screen.load_image("Jeep.png")
# Wait until touch sensor pressed, speak
while not sensor_touch.pressed():
    pass

# Set a stopwatch to delay colour detection for 2 seconds
stopwatch = StopWatch()
time_delay = TIME_WAIT * 2000

# Set up comms to listen for dinosaur sound commands
host = '0.0.0.0'
port = 12346  # new port for receiving sound commands

r = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
r.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
r.bind((host, port))
r.listen(5)
r.settimeout(0.5)
msg = "BLANK"
print("Listening for sound commands...")

# Define thread to listen for sounds
#def listen_for_sounds():
#    while True:
#        if not playing_sound:
#            receive_message()

# Define thread to look for zones
def look_for_zones():
    #global looking_for_zones
    #while looking_for_zones:
    while True:
        speak_message()
        if not playing_sound:
            receive_message()

# Start thread
#_thread.start_new_thread(listen_for_sounds, ())
_thread.start_new_thread(look_for_zones, ())

# Wait until gate opens
while sensor_us.distance() < THRESHOLD_OBSTACLE:
    pass

zone = "START"

while trex_not_detected:

    detected_colour = None
    # Add some delay to avoid immediate / incorrect colour detection
    if stopwatch.time() < time_delay:
        drive()
        continue

    else:
        # Detect colours
        detected_colour = detect_colours()
        # If appropriate colour detected, add delay back
        if detected_colour != 'Unknown':
            #print('Detected colour:', detected_colour)
            stopwatch.reset()

        drive()

# Check for T-Rex, avoid obstacle
drive_base.stop()
while sensor_us.distance() > THRESHOLD_OBSTACLE:
    pass

#looking_for_zones = False
avoid_obstacle()
zone = "END"

# Continue listening after 'shutting' down
while True:
    receive_message()