#!/usr/bin/env python3.6

from controller import Robot, Motor, PositionSensor, InertialUnit, Lidar, Display
import numpy as np
#PLOTTING FUNCTIONS
class Grapher:
    def __init__(self, display):
        '''
        display: The display you wish to use with this.
        '''
        self.display = display 
        self.width = display.getWidth()
        self.height = display.getHeight()
        self.defaultPointSize = int((self.height + self.width)/200)
    def drawPointCorner(self, x, y, size=None, color=0x00FFFF):
        '''
        Draws square points wrt the top right corner. The size is in pixels.
        '''
        if(size is None):
            size = self.defaultPointSize
        self.display.setColor(color)
        for row in range(max(0, y - size//2), min(self.height, y + size//2)):
            for column in range(max(0, x - size//2), min(self.width, x + size//2)):
                self.display.drawPixel(column, row)
    def drawPointCenter(self, x, y, size=None, color=0x00FFFF):
        '''
        Draws square points taking the center (width//2, height//2) to be 0, 0.
        '''
        self.drawPointCorner(x + self.width//2, self.height//2-y, size, color)
    def drawPointsListCenter(self, x_s, y_s, size = None, color=0x00FFFF):
        '''
        Draws points from two lists. If len_x != len_y, the trailing of x or y are truncated.        
        '''
        length = min(len(x_s), len(y_s))
        for i in range(length):
            self.drawPointCenter(x_s[i], y_s[i], size, color)

    def clear(self, color=0x000000):
        self.display.setColor(color)
        self.display.fillRectangle(0, 0, self.width, self.height)

    def lidar_plot(self, x_data, y_data, DISPLAY_SCALING_FACTOR, color=0x00FFFF):
        self.clear()
        self.lidar_plot_without_clearing(x_data, y_data, DISPLAY_SCALING_FACTOR, color)
    def lidar_plot_without_clearing(self, x_data, y_data, DISPLAY_SCALING_FACTOR, color=0x00FFFF):
        x_s = []
        y_s = []
        for i in range(len(x_data)):
            x_s.append(int(x_data[i]*DISPLAY_SCALING_FACTOR))
            y_s.append(int(y_data[i]*DISPLAY_SCALING_FACTOR))
        # self.drawPointCenter(0, 0, color=0xFF0000)
        self.drawPointsListCenter(x_s, y_s, size=5,color=color)
def lidar_filter(imageArray, SIZES,RANGES,EPSILON):
    theta_data = np.zeros((SIZES[1],)).tolist()
    for layer in range(SIZES[0]):
        for theta in range(SIZES[1]):
                point_range = imageArray[layer][theta]
                if(point_range < RANGES[layer]*(EPSILON)):
                    theta_data[theta] = point_range
    return theta_data







#INITIALIZATION CODE HERE.
VELOCITY = 30
WHEEL_RADIUS = 0.0075
#ROBOT INITIALIZATION
robot = Robot()
timestep = int(robot.getBasicTimeStep())
#LIDAR AND PLOTTING INITIALIZATION
SIZES = (16, 512)
ranges_str = "1.13114178 0.85820043 0.57785118 0.43461093 0.38639969 0.31585345 0.2667459 0.23062678 0.21593061 0.19141567 0.17178488 0.15571462 0.14872716 0.13643947 0.12597121 0.11696267"
RANGES = [float(i) for i in ranges_str.split(' ')]
EPSILON = 0.6
DISPLAY_SIZE = (1024, 1024)
DISPLAY_SCALING_FACTOR = 0.9*1024/5
PLOT_UPDATE_RATE = 1
lidar = robot.getLidar("lidar")
lidar.enable(timestep)
i = 0
display = robot.getDisplay("extra_display")
grapher = Grapher(display)



#MOTOR AND ENCODER INITIALIZATION
left_motor = robot.getMotor("left motor")
right_motor = robot.getMotor("right motor")

left_motor_sensor = left_motor.getPositionSensor()
right_motor_sensor = right_motor.getPositionSensor()

left_motor_sensor.enable(timestep)
right_motor_sensor.enable(timestep)

left_motor.setPosition(float('inf'))
right_motor.setPosition(float('inf'))

#KEYBOARD MESSAGE
print("Keyboard control has been disabled.")

#INERTIAL UNIT INITIALIZATION
inertial_unit = robot.getInertialUnit("inertial unit")
inertial_unit.enable(timestep)

#ODOMETRY INITIALIZATION
_, _, alpha = inertial_unit.getRollPitchYaw()
x = 0 
z = 0
DTHETA_L = 0
DTHETA_R = 0
THETA_L_PREV = left_motor_sensor.getValue()
THETA_R_PREV = right_motor_sensor.getValue()
distance = 0

#WALL-FOLLOWING BEHAVIOUR INITIALIZATION
#ENSURE THAT LEFT + SPREAD < 512 

CENTER_DELTA = 0.4
LEFT_DELTA = 0.7
FRONT = 0
RIGHT = 128
BACK = 256
LEFT = 384
LEFT_SPREAD_LOW = -30
LEFT_SPREAD_HIGH = 100
CENTER_SPREAD = 50
turnLeft = True
goForward = False
LAG = 50
j = LAG
while robot.step(timestep) != -1:
        #HANDLING INERTIAL UNIT BEHAVIOUR -- INCOMPLETE!
        _, _, alpha = inertial_unit.getRollPitchYaw()
        DTHETA_L = left_motor_sensor.getValue() - THETA_L_PREV
        DTHETA_R = right_motor_sensor.getValue() - THETA_R_PREV
        THETA_L_PREV += DTHETA_L
        THETA_R_PREV += DTHETA_R
        if(goForward):
            DTHETA = (DTHETA_L + DTHETA_R)/2
            x += WHEEL_RADIUS*DTHETA*np.sin(alpha)
            z += WHEEL_RADIUS*DTHETA*np.cos(alpha)
            distance += abs(WHEEL_RADIUS*DTHETA)

        #LIDAR SENSOR HANDLING PART 1 for WALL-FOLLOWING BEHAVIOUR
        imageArray = np.array(lidar.getRangeImageArray()).T
        theta_data = lidar_filter(imageArray, SIZES,RANGES,EPSILON)
        theta_data_aligned = []
        for theta in range(len(theta_data)):
            new_theta =(-np.pi/2 -  2*np.pi*theta/512 + alpha)
            if(theta_data[theta] != 0):  
                theta_data_aligned.append((new_theta, theta_data[theta]))
        x_data = []
        y_data = []
        for i in range(len(theta_data_aligned)):
            x_data.append(theta_data_aligned[i][1]*np.cos(theta_data_aligned[i][0]))
            y_data.append(theta_data_aligned[i][1]*np.sin(theta_data_aligned[i][0]))

        for i in range(len(x_data)):
            x_data[i] += x 
            y_data[i] -= z
        grapher.lidar_plot_without_clearing([x],[-z],DISPLAY_SCALING_FACTOR,color=0xFF0000)
        if(goForward and not turnLeft and j > LAG):
            grapher.lidar_plot_without_clearing(x_data, y_data, DISPLAY_SCALING_FACTOR)  
        elif(goForward and not turnLeft):
            j += 1
        #WALL-FOLLOWING BEHAVIOUR
        turnLeft = True
        goForward = False
        for i in range(LEFT-LEFT_SPREAD_LOW,LEFT+LEFT_SPREAD_HIGH+1):
            if(theta_data[i] != 0 and theta_data[i] < LEFT_DELTA):
                turnLeft = False
        if(turnLeft):
            j = 0
            left_motor.setVelocity(-VELOCITY)
            right_motor.setVelocity(VELOCITY)
        else:
            goForward = True
            for i in range(FRONT - CENTER_SPREAD , FRONT + CENTER_SPREAD + 1):
                if(theta_data[i] != 0 and theta_data[i] < CENTER_DELTA):
                    goForward = False
            if(goForward):
                left_motor.setVelocity(VELOCITY)
                right_motor.setVelocity(VELOCITY)
            else:
                j = 0
                left_motor.setVelocity(VELOCITY)
                right_motor.setVelocity(-VELOCITY)
        
        