import vrep
import sys
import time
import numpy as np
from PID import PID


class Robot:
    def __init__(self):
        vrep.simxFinish(-1)

        self.timeSleep = 0.1

        self.distDefault = 0.576
        self.speed = 0.2

        self.pid = PID(0.4, 0.4, 0)

        self.clientId = vrep.simxStart('127.0.0.1', 19999, True, True, 5000, 5)
        assert self.clientId != -1, 'Failed connecting to remote API server'

        e = vrep.simxStartSimulation(self.clientId, vrep.simx_opmode_oneshot_wait)
        self.checkErrorCode(e, "Unable to start simulation", True)

        e, self.leftMotorHandle = vrep.simxGetObjectHandle(self.clientId, 'Pioneer_p3dx_leftMotor', vrep.simx_opmode_oneshot_wait)
        self.checkErrorCode(e, "Can not find left motor", True)

        e, self.rightMotorHandle = vrep.simxGetObjectHandle(self.clientId, 'Pioneer_p3dx_rightMotor', vrep.simx_opmode_oneshot_wait)
        self.checkErrorCode(e, "Can not find right motor", True)

        e, self.LeftSensorHandle = vrep.simxGetObjectHandle(self.clientId, 'Proximity_sensor',
                                                            vrep.simx_opmode_oneshot_wait)
        self.checkErrorCode(e, "Can not find left ray sensor", True)

        e, self.TopSensorHandle = vrep.simxGetObjectHandle(self.clientId, 'Proximity_sensor0', vrep.simx_opmode_oneshot_wait)
        self.checkErrorCode(e, "Can not find top ray sensor", True)

        sec, msecond = vrep.simxGetPingTime(self.clientId)
        print("Ping time: %f" % (sec + msecond / 1000.0))
        print("Initialization finished")

    @staticmethod
    def checkErrorCode(errorCode, message, need_exit=False):
        if errorCode != vrep.simx_return_ok:
            print("ERROR: Code {}. {}".format(errorCode, message))
            if need_exit:
                sys.exit()

    def setMotorSpeed(self, leftSpeed, rightSpeed):
        e = vrep.simxSetJointTargetVelocity(self.clientId, self.leftMotorHandle, leftSpeed,
                                            vrep.simx_opmode_oneshot_wait)
        self.checkErrorCode(e, "SetJointTargetVelocity for left motor got error code")
        e = vrep.simxSetJointTargetVelocity(self.clientId, self.rightMotorHandle, rightSpeed,
                                            vrep.simx_opmode_oneshot_wait)
        self.checkErrorCode(e, "SetJointTargetVelocity for right motor got error code")

        print("Motor speed set to {} {}".format(leftSpeed, rightSpeed))

    def getSensorData(self, sensor):
        (e, detectionState, detectedPoint, detectedObjectHandle,
         detectedSurfaceNormalVector) = vrep.simxReadProximitySensor(self.clientId, sensor, vrep.simx_opmode_buffer)
        self.checkErrorCode(e, "simxReadProximitySensor error")

        return detectionState, np.linalg.norm(detectedPoint), detectedSurfaceNormalVector

    def getPidCorrection(self, dist):
        errorThree = self.distDefault - round(dist, 3)
        error = errorThree

        correction = self.pid(error)

        return errorThree, correction

    def startSimulation(self):
        vrep.simxReadProximitySensor(self.clientId, self.LeftSensorHandle, vrep.simx_opmode_streaming)
        vrep.simxReadProximitySensor(self.clientId, self.TopSensorHandle, vrep.simx_opmode_streaming)

        while vrep.simxGetConnectionId(self.clientId) != -1:
            isDetected, distance, vector = self.getSensorData(self.LeftSensorHandle)
            isDetectedTop, distance2, vector2 = self.getSensorData(self.TopSensorHandle)

            if isDetectedTop:
                errorDistanse, pidCorrect = self.getPidCorrection(distance2)
                self.setMotorSpeed(self.speed - pidCorrect, self.speed + pidCorrect)
                print("distance: {} {} | pid: {} | error: {}".format(distance, vector, pidCorrect, errorDistanse))
            else:
                self.setMotorSpeed(self.speed + 0.05, self.speed - 0.05)

            time.sleep(self.timeSleep)
        vrep.simxFinish(self.clientId)


if __name__ == '__main__':
    robot = Robot()
    robot.startSimulation()