from IPython.display import SVG
SVG(filename='n-pendulum-with-cart.svg')

from sympy import symbols, trigsimp
from sympy.physics.mechanics import *

# Initialize the pretty-printing system
init_vprinting()

# Number of Links
n = 3

# Generalized Coordinates and Speeds
q = dynamicsymbols('q:' + str(n + 1))
u = dynamicsymbols('u:' + str(n + 1))

# Input Force
f = dynamicsymbols('f')

# Constants:
m = symbols('m:' + str(n + 1))         # Mass of each bob
l = symbols('l:' + str(n))             # Length of each link
g, t = symbols('g t')                  # Gravity and time

I = ReferenceFrame('I')
O = Point('O')
O.set_vel(I, 0)

P0 = O.locatenew('P0', q[0] * I.x)
P0.set_vel(I, u[0] * I.x)
Pa0 = Particle('Pa0', P0, m[0])

frames = [I]
points = [P0]
particles = [Pa0]
forces = [(P0, f * I.x - m[0] * g * I.y)]
kindiffs = [q[0].diff() - u[0]]

for i in range(n):
    # Create a new frame oriented with current q and u
    Bi = I.orientnew('B' + str(i), 'Axis', [q[i + 1], I.z])
    Bi.set_ang_vel(I, u[i + 1] * I.z)
    frames.append(Bi)

    # Create a new point at the origin of the frame
    Pi = points[-1].locatenew('P' + str(i + 1), l[i] * Bi.x)
    Pi.v2pt_theory(points[-1], I, Bi)
    points.append(Pi)
    
    # Create a particle at that point
    Pai = Particle('Pa' + str(i + 1), Pi, m[i + 1])
    particles.append(Pai)

    # Add gravity acting at that point
    forces.append((Pi, -m[i + 1] * g * I.y))

    # Define the kinematic equations for this point
    kindiffs.append(q[i + 1].diff(t) - u[i + 1])

kane = KanesMethod(I, q_ind=q, u_ind=u, kd_eqs=kindiffs)
fr, frstar = kane.kanes_equations(forces, particles)

trigsimp(fr)

trigsimp(frstar)

from pydy.system import System
from numpy import hstack, zeros, linspace, pi, ones
from matplotlib import pyplot as plt
%matplotlib inline

arm_length = 1. / n
bob_mass = 0.01 / n

# Substitution dictionary, for applying the parameters
parameter_dict = {g: 9.81}
parameter_dict.update({i: arm_length for i in l})
parameter_dict.update({i: bob_mass for i in m})

print(parameter_dict)

sys = System(kane, constants=parameter_dict)

x0 = hstack(( 0, pi / 2 * ones(len(q) - 1), 1e-3 * ones(len(u))))
t = linspace(0, 10, 1000)

sys.initial_conditions = dict(zip(q + u, x0))
sys.times = t
y = sys.integrate()

plt.plot(t, y[:, :y.shape[1] / 2])
plt.xlabel('Time [sec]')
plt.legend(q)
plt.title("Generalized Coordinates");

plt.plot(t, y[:, y.shape[1] / 2:])
plt.xlabel('Time [sec]')
plt.legend(u)
plt.title("Generalized Speeds");

from util import animate_pendulum

%%capture
animate_pendulum(t, y, arm_length, filename="open-loop.mp4")

%%capture
!mplayer open-loop.mp4

from control import ctrb, lqr
from numpy.linalg import matrix_rank, eigvals
from sympy import matrix2numpy

equilibrium_point = hstack(( 0, pi / 2 * ones(len(q) - 1), zeros(len(u)) ))

# Create a substitution dictionary to apply the point
equilibrium_dict = dict(zip(q + u, equilibrium_point))
equilibrium_dict.update({i.diff(): 0 for i in u})

A, B, inputs = kane.linearize(op_point=(parameter_dict, equilibrium_dict), A_and_B=True, new_method=True)
A = matrix2numpy(A, float)
B = matrix2numpy(B, float)

eigvals(A)

matrix_rank(ctrb(A, B)) == A.shape[0]

K, X, E = lqr(A, B, ones(A.shape), 1)

sys.specifieds = {f: lambda x, t: K.dot(equilibrium_point - x)}
x0 = hstack(( 0, pi / 2 * ones(len(q) - 1), 1 * ones(len(u)) ))
sys.initial_conditions = dict(zip(q + u, x0))

y = sys.integrate()

plt.plot(t, y[:, :y.shape[1] / 2])
plt.xlabel('Time [sec]')
plt.legend(q)
plt.title("Generalized Coordinates");

plt.plot(t, y[:, y.shape[1] / 2:])
plt.xlabel('Time [sec]')
plt.legend(u)
plt.title("Generalized Speeds");

%%capture
animate_pendulum(t, y, arm_length, filename="closed-loop.mp4")

%%capture
!mplayer closed-loop.mp4

from sympy.printing import ccode
from sympy.utilities.codegen import codegen
from sympy.utilities.autowrap import autowrap
from sympy import sin, sqrt, lambdify

a, b, c = symbols('a, b, c')
expr = a + b*c + sin(a) + sqrt(b)
expr

ccode(expr)

code, header = codegen(("test", expr), "c")

print(header[1])

print(code[1])

f_native = autowrap(expr)

f_native(1, 2, 3)

f_python = lambdify((a, b, c), expr)

f_python(1, 2, 3)

%timeit f_native(1, 2, 3)

%timeit f_python(1, 2, 3)