import pylab
import math

GRAVITY = 9.81

class Pendolo:

    def __init__(self, _M, _b):
        self.w = 0
        self.theta = 0
        self.M = _M
        self.b = _b

    def evaluate(self, _input, delta_t):
        theta_temp = self.theta + delta_t * self.w
        w_temp = -GRAVITY * delta_t * self.theta + \
          (1 - self.b * delta_t/self.M) * self.w + \
            delta_t / self.M * _input

        self.theta = theta_temp;
        self.w = w_temp;



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

    def __init__(self, kp, ki, kd, sat):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.saturation = sat
        self.integral = 0
        self.prev_error = 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
        deriv = (error - self.prev_error) / delta_t
        self.prev_error = error
        output = self.kp * error + self.ki * self.integral + self.kd * deriv
        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

t = 0.0

braccio = Pendolo(6.0, 4.0)
#speed_controller = PISat(2, 10, 120)
speed_controller = PIDSat(10, 30, 20, 120)

w_target = math.radians(30)

vettore_theta = [ ]
vettore_w = [ ]
vettore_tempi = [ ]
vettore_output = [ ]

while t < 20:

    output = speed_controller.evaluate(w_target, braccio.theta, delta_t)
    braccio.evaluate(output, delta_t)

    t = t + delta_t

    vettore_output.append(output)
    vettore_w.append(braccio.w)
    vettore_theta.append(math.degrees(braccio.theta))
    vettore_tempi.append(t)


pylab.figure(1)
pylab.plot(vettore_tempi, vettore_w, 'r-+', label='vel, w(t)')
pylab.xlabel('time')
pylab.legend()

pylab.figure(2)
pylab.plot(vettore_tempi, vettore_theta, 'b-+',
		label='position, theta(t)')
pylab.xlabel('time')
pylab.legend()

pylab.figure(3)
pylab.plot(vettore_tempi, vettore_output, 'b-+',
		label='output controller')
pylab.xlabel('time')
pylab.legend()

pylab.show()