import pylab

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 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


delta_t = 1e-3 # 1 ms

robot = Massa(6.0, 25.0)
ctrl = PISat(150, 400, 200)

target_speed = 5  # 5 metri/s

t = 0.0
vettore_target = [ ]
vettore_vel = [ ]
vettore_tempi = [ ]
vettore_f = [ ]

while t < 2:
    current_speed = robot.get_speed()

    f = ctrl.evaluate(target_speed, current_speed, delta_t)

    robot.evaluate(f, delta_t)

    t = t + delta_t
    vettore_f.append(f)
    vettore_vel.append(robot.get_speed())
    vettore_target.append(target_speed)
    vettore_tempi.append(t)


pylab.figure(1)
pylab.plot(vettore_tempi, vettore_vel, 'r-+', label='vel, v(t)')
pylab.plot(vettore_tempi, vettore_target, 'b-+', label='target vel')
pylab.xlabel('time')
pylab.legend()

pylab.figure(3)
pylab.plot(vettore_tempi, vettore_f, 'k-+', label='force, f(t)')
pylab.xlabel('time')
pylab.legend()

pylab.show()