import pylab
import math

class PISat:

    def __init__(self, kp, ki, sat):
        self.kp = kp
        self.ki = ki
        self.saturation = sat
        self.integral = 0
        self.saturation_flag = False

    def evaluate(self, target, current, delta_t):
        error = target - current
        if not(self.saturation_flag):
            self.integral = self.integral + error * delta_t
        output = self.kp * error + self.ki * self.integral
        if output > self.saturation:
            output = self.saturation
            self.saturation_flag = True
        elif output < -self.saturation:
            output = -self.saturation
            self.saturation_flag = True
        else:
            self.saturation_flag = False
        return output


class Massa:

    def __init__(self, _M, _b):
        self.M = _M
        self.b = _b
        # variabili di stato, p = posizione, v = velocita'
        self.p = 0
        self.v = 0

    def evaluate(self, _input, dt):
        self.p = self.p + delta_t * self.v
        self.v = (1 - self.b * dt/self.M) * self.v + dt / self.M * _input

    def get_position(self):
        return self.p

    def get_speed(self):
        return self.v


class Robot:

    def __init__(self, _M, _b, _wheelbase):
        self.__left_w = Massa(_M / 2.0, _b)
        self.__right_w = Massa(_M / 2.0, _b)
        self.wheelbase = _wheelbase
        self.x = 0
        self.y = 0
        self.theta = 0
        # posizioni precedenti (servono per calcolare l'odometria)
        self.__prev_pL = 0
        self.__prev_pR = 0
        # controllori velocita'
        self.__speed_pi_left = PISat(150,400,200)
        self.__speed_pi_right = PISat(150,400,200)
        # velocita' attuali
        self.current_vl = 0
        self.current_vr = 0

    def evaluate(self, vl, vr, delta_t):
        # aggiornamento odometria
        pL = self.__left_w.get_position()
        pR = self.__right_w.get_position()

        delta_L = pL - self.__prev_pL
        delta_R = pR - self.__prev_pR

        self.__prev_pL = pL
        self.__prev_pR = pR

        delta_theta = (delta_R - delta_L) / self.wheelbase
        delta_linear = (delta_L + delta_R) / 2.0

        self.x = self.x + delta_linear * math.cos(self.theta + delta_theta / 2.0)
        self.y = self.y + delta_linear * math.sin(self.theta + delta_theta / 2.0)
        self.theta = self.theta + delta_theta

        if self.theta > math.pi:
            self.theta = self.theta - 2*math.pi
        if self.theta < -math.pi:
            self.theta = 2*math.pi + self.theta

        # calcolo controllori di velocita' ruote
        self.current_vl = self.__left_w.get_speed()
        self.current_vr = self.__right_w.get_speed()
        output_L = self.__speed_pi_left.evaluate(vl, self.current_vl, delta_t)
        output_R = self.__speed_pi_right.evaluate(vr, self.current_vr , delta_t)

        # applicazione dell'output ai motori
        self.__left_w.evaluate(output_L, delta_t)
        self.__right_w.evaluate(output_R, delta_t)



class PolarControl:

    def __init__(self, robot, kp_lin, sat_lin, kp_angular, sat_angular):
        self.kp_lin = kp_lin
        self.sat_lin = sat_lin
        self.kp_angular = kp_angular
        self.sat_angular = sat_angular
        self.robot = robot

    def evaluate(self, target_x, target_y):

        dx = target_x - self.robot.x
        dy = target_y - self.robot.y

        target_heading = math.atan2(dy, dx)
        distance = math.hypot(dx, dy)

        heading_error = target_heading - robot.theta

        v = self.kp_lin * distance
        w = self.kp_angular * heading_error

        if v > self.sat_lin:
            v = self.sat_lin
        elif v < -self.sat_lin:
            v = - self.sat_lin

        if w > self.sat_angular:
            w = self.sat_angular
        elif w < -self.sat_angular:
            w = - self.sat_angular

        vl = v - (w * self.robot.wheelbase / 2)
        vr = v + (w * self.robot.wheelbase / 2)
        return (vl, vr)


delta_t = 0.001
robot = Robot(6.0, 25.0, 0.3)

t = 0
time_array = [ ]
x_array = [ ]
y_array = [ ]
theta = [ ]
vl_array = [ ]
vr_array = [ ]

p = PolarControl(robot,
                 1, 1.5,
                 5, 4)

while t < 15:

    (vl, vr) = p.evaluate(1, 10)

    robot.evaluate(vl, vr, delta_t)

    time_array.append(t)
    x_array.append(robot.x)
    y_array.append(robot.y)
    vl_array.append(robot.current_vl)
    vr_array.append(robot.current_vr)

    t = t + delta_t



pylab.figure(1)
pylab.plot(x_array, y_array, 'r-+', label='posizione robot')
pylab.legend()

pylab.figure(2)
pylab.plot(time_array, vl_array, 'r-+', label='vl')
pylab.plot(time_array, vr_array, 'b-+', label='vr')
pylab.legend()

pylab.show()