#!/usr/bin/env python3

from time import sleep
from ev3dev2.motor import Motor, MoveTank, SpeedPercent, OUTPUT_A, OUTPUT_B, OUTPUT_C, OUTPUT_D 
from ev3dev2.sensor import INPUT_2,INPUT_3,INPUT_4
from ev3dev2.sensor.lego import ColorSensor, UltrasonicSensor, GyroSensor
from ev3dev2.led import Leds
from pixycamev3.pixy2 import Pixy2

class Object_Sorter(): 
    # Define type of global variables
    state: str
    last_state: str
    us: UltrasonicSensor
    cs: ColorSensor
    gyro: GyroSensor
    m_left: Motor
    m_right: Motor
    tank: MoveTank
    leds: Leds
    pixy2: Pixy2
    params: dict
    states: dict
    cycles: int

    # Define all inputs, outputs of the Robot and initialize
    def __init__(self, **args):
        # Camera and sensors are inputs
        self.us = UltrasonicSensor(args.get("ultrasonic_input",INPUT_2))
        self.cs = ColorSensor(args.get("color_input",INPUT_3))
        self.gyro = GyroSensor(args.get("gyro_input",INPUT_4))
        self.pixy2 = Pixy2(port=args.get("camera_port",1), i2c_address=0x54)
        # Motors are outputs
        self.m_left = Motor(args.get("leftmotor_output",OUTPUT_A)) 
        self.m_right = Motor(args.get("rightmotor_output",OUTPUT_D))
        # Also define drive tank to move syncronized
        self.tank = MoveTank(args.get("leftmotor_output",OUTPUT_A), args.get("rightmotor_output",OUTPUT_D)) 
        self.leds = Leds()

        # The motor is initialized to run counter-clockwise (gegen den Uhrzeigersinn)
                # use 'normal' to initialize it to run clockwise (im Uhrzeigersinn)
        self.m_left.polarity = args.get("leftmotor_polarity",'inversed')
        self.m_right.polarity = args.get("rightmotor_polarity",'inversed') 
        
        # Read inputs for the given experiment / enviroment
        args["sig_let"] = args.get("signature_dice_remains",1) # Signatur of orange dice
        args["sig_sort"] = args.get("signature_dice_sorted",2)# Signatur of green dice
        args["fast"] = args.get("speed_drive", 30) # driving speed
        args["slow"] = args.get("speed_search", 10) # searching speed
        args["dist_us"] = args.get("distance_us",12) # detection range ultrasonic sensor
        args["line"] = args.get("max_reflection_outline", 10) # max. reflection of outline 
        
        # Safe Parameters for later, can be changed
        self.params = args

        # initalization state defined
        self.state = "search" # Initiate search state for beginning
        self.last_state = "search" # Variable for saving last state 
        self.cycles = 0 # control variable
        self.gyro.reset() # reset gyro sensor

        self.states = { # dictionary of the states/classes 
            "search": self.search,
            "get": self.get,
            "sort": self.sort,
            "avoid": self.avoid,
            "edge": self.edge
        }

    # Define method/state search
    def search (self): 
        # turn slowly aroud on place and search for blocks, 
        # if nothing is found change searchplace
        self.leds.set_color("LEFT", "BLACK")
        self.leds.set_color("RIGHT", "BLACK")
        # turn on place to the right
        self.m_right.on(SpeedPercent(-1*self.params["slow"]))
        self.m_left.on(SpeedPercent(self.params["slow"])) 
        # search for blocks
        nr_sort, sort_block = self.pixy2.get_blocks(self.params["sig_sort"], 2)
        nr_let, let_block = self.pixy2.get_blocks(self.params["sig_let"], 2)
        if self.cs.reflected_light_intensity < self.params["line"]: # if outline is reached change state to edge
            self.tank.stop()
            self.last_state = self.state
            self.state = "edge"
        elif self.us.distance_centimeters < self.params["dist_us"]: # if a block is currently in the arms change to sort
            self.last_state = self.state 
            self.state = "sort"
        elif nr_sort > 0: # when block to sort is found then change to get 
            self.tank.stop()
            self.last_state = self.state
            self.state = "get"
        elif nr_let > 0: # when block that remains is found then change to avoid 
            self.last_state = self.state
            self.state = "avoid"
        elif self.gyro.angle > 720: # when the robot does not find a block drive a little bit forward
            rounds=0 # rounds of driving forward
            while self.cs.reflected_light_intensity > self.params["line"] and rounds <= 5:
                self.tank.on(self.params["fast"],self.params["fast"])
                sleep(0.2)
                rounds+=1 
            if self.cs.reflected_light_intensity <= self.params["line"] : 
                self.tank.stop()
                self.last_state = self.state
                self.state = "edge"
            self.gyro.reset()
            rounds=0 # Reset control variable

    # Define method/state get  
    def get (self): 
        # drive twoards block and get it into arms
        self.leds.set_color("LEFT", "GREEN")
        self.leds.set_color("RIGHT", "BLACK")
        # drive straight ahead
        self.tank.on(self.params["fast"],self.params["fast"])
        # find the block for sorting again
        nr_sort, sort_block = self.pixy2.get_blocks(self.params["sig_sort"], 2)
        if nr_sort > 0: # if block is found again get position
            x = sort_block[0].x_center #get where in the field of vision the block lays
            self.cycles = 0
        else: # when block is not found, virtually place the robot in the middle of the vision 
            x = 130 
            self.cycles +=1
            if 4 <= self.cycles < 6: # camera has not found block for 4 cycles
                self.tank.on_for_seconds(-1*self.params["fast"],-1*self.params["fast"],0.4) # drive backwards max. 2 times
            elif 6 == self.cycles: # if block is not found drive a little bit to left
                self.m_right.on(SpeedPercent(self.params["slow"]))
                self.m_left.on(SpeedPercent(-1*self.params["slow"])) 
                sleep(0.2)
            elif self.cycles > 6: # when block is still not found go back to search
                self.cycles = 0
                self.state = "search"   
                self.gyro.reset()
        # Decide what robot should do with the gathered information
        if self.cs.reflected_light_intensity < self.params["line"]: # robot detected line change to edge
            self.tank.stop()
            self.last_state = self.state
            self.state = "edge" # When colorsensor detects line then switch state 
        elif self.us.distance_centimeters <= self.params["dist_us"]: # block is in the arms of robot change to sort
            self.last_state = self.state
            self.state = "sort"  # when block distance is small switch state 
        else: # correct drive direction to get block into arms
            # when block on left side drive more left
            if x < 110: # turn left
                self.m_right.on(SpeedPercent(self.params["slow"]))
                self.m_left.on(SpeedPercent(-1*self.params["slow"]))  
            # when block on right side drive more right
            elif x > 150: # turn right
                self.m_right.on(SpeedPercent(-1*self.params["slow"]))
                self.m_left.on(SpeedPercent(self.params["slow"]))   
        

    # Define method/state sort
    def sort (self):  
        # block is caught, drive to the edge
        self.leds.set_color("LEFT", "GREEN")
        self.leds.set_color("RIGHT", "GREEN")
        # drive straight ahead
        self.tank.on(self.params["fast"],self.params["fast"])
        # search for block to let be in area
        nr_let, let_block = self.pixy2.get_blocks(sig_let, 2)
        if self.cs.reflected_light_intensity < self.params["line"]: # when reached line change to edge
            self.tank.stop()
            self.last_state = self.state
            self.state = "edge" 
        elif nr_let > 0: 
            self.last_state = self.state
            self.state = "avoid"

    # Define method/state avoid
    def avoid (self):
        # found green block, do not want to kick it out 
        self.leds.set_color("LEFT", "ORANGE")
        self.leds.set_color("RIGHT", "ORANGE")
        # search for green block again
        nr_let, let_block = self.pixy2.get_blocks(sig_let, 2)
        if nr_let > 0: # if block is found again get position
            x = let_block[0].x_center # get where in the field of vision the block lays
        else: # when block is not found, virtually place the robot in the middle of the vision 
            x = 130
        # When block lays in middle of the vision drive around it 
        if 100 < x < 150: 
            self.m_right.on(SpeedPercent(-1*self.params["fast"]))
            self.m_left.on(SpeedPercent(self.params["fast"]))
            sleep(0.2)
        self.state = self.last_state
        self.last_state = "avoid"

    # Define method/state edge
    def edge (self):
        # robot reached outerline and should turn around
        self.leds.set_color("LEFT", "BLACK")
        self.leds.set_color("RIGHT", "RED")
        self.gyro.reset() # reset gyro sensor 
        # drive straight backwards
        self.tank.on_for_rotations(-1*self.params["fast"],-1*self.params["fast"],2)
        while self.gyro.angle < 100: #turn around for 100 degrees
            self.m_right.on(SpeedPercent(-1*self.params["fast"])) 
            self.m_left.on(SpeedPercent(self.params["fast"]))
        self.last_state = self.state
        self.state = "search"
        self.gyro.reset() # reset gyro sensor again
    
    def get_params(self):
        return self.params

    def run(self): 
        print('start')
        while True: # endless-loop for switching states        
            func = self.states.get(self.state) # function that gets state 
            func() # activate current state
            sleep(0.3) # wait for 0.3 seconds between the different states 

if __name__ == "__main__":
    # All arguments that can be changed are given here
    arguments = dict(
        #ultrasonic_input = INPUT_2,
        #color_input = INPUT_3,
        #gyro_input = INPUT_4,
        #camera_port = 1,
        #leftmotor_output = OUTPUT_A,
        #rightmotor_output = OUTPUT_D,
        #leftmotor_polarity = 'inveresed',
        #rightmotor_polarity = 'inveresed',
        #signature_dice_remains = 1,   
        #signature_dice_sorted = 2,
        #speed_drive = 30,
        #speed_search = 10,
        #distance_us =10,
        max_reflection_outline = 10

    )
    
    os = Object_Sorter(**arguments)
    os.run()
    