#
# pendolo.py
#
import pylab

G = 9.81

class RoboticArm:

    def __init__(self, m, b):
        self.__mass = m
        self.__friction = b
        self.__theta = 0.0 # condizione iniziale
        self.__omega = 0.0 # condizione iniziale

    def evaluate(self, _input, delta_t):
        new_theta = self.__theta + delta_t * self.__omega
        new_omega = -G * delta_t * self.__theta + \
          (1 - self.__friction * delta_t / self.__mass) * self.__omega + \
          delta_t/self.__mass * _input

        self.__theta = new_theta
        self.__omega = new_omega

    def get_theta(self):
        return self.__theta

    def get_omega(self):
        return self.__omega


arm = RoboticArm(6.0, 5.0)
delta_t = 1e-3 # 1ms

duration = 20.0 # 20 seconds of simulations

ticks = int(duration / delta_t)

f_k = 20 # newton of constant input

_input = [ ]
thetas = [ ]
omegas = [ ]
times = [ ]

for i in range(0, ticks):
    _input.append(f_k)
    arm.evaluate(f_k, delta_t)
    thetas.append(arm.get_theta())
    omegas.append(arm.get_omega())
    times.append(i * delta_t)
    f_k = 0

pylab.plot(times, omegas, 'r-+', label='speed, w(t)')
#pylab.plot(times, thetas, 'b-+', label='position, theta(t)')
pylab.xlabel('time')
pylab.legend()
pylab.show()