#!/usr/bin/env python
# coding: utf-8

# # Dealing with configuration constraints that can't be solved analytically
# 
# This is example is the classic 4 bar linkage. In this case there are two particles at P2 and P3 with mass.
# 
# ![](https://objects-us-east-1.dream.io/mae223/2019f/four-bar-linkage-sketch.png)

# In[1]:


import sympy as sm
import sympy.physics.mechanics as me
me.init_vprinting()


# We have three non-minimal coordinates but only one is needed to specify the configuration.

# In[2]:


q1, q2, q3 = me.dynamicsymbols('q1, q2, q3')
u1 = me.dynamicsymbols('u1')
T = me.dynamicsymbols('T')


# In[3]:


l1, l2, l3, l4 = sm.symbols('l1:5')
m2, m3, g = sm.symbols('m2, m3, g')


# In[4]:


N = me.ReferenceFrame('N')
A = N.orientnew('A', 'Axis', (q1, N.z))
B = A.orientnew('B', 'Axis', (q2, -A.z))
C = B.orientnew('C', 'Axis', (q3, -B.z))


# In[5]:


P1 = me.Point('P1')
P2 = P1.locatenew('P2', l1*A.x)
P3 = P2.locatenew('P3', l2*B.x)
P4 = P3.locatenew('P4', l3*C.x)


# This is the vector equation that will be used to create the configuration constraints.

# In[6]:


loop = P4.pos_from(P1) - l4*N.x
loop


# There are two configuration constraints and thus only one of the configuration variables is necessary for speciying the complete coniguration of the system. In this case it may be possible to solve these two nonlinear equations for q2 and q3, but that is often not the case.

# In[7]:


config_con1 = loop.dot(N.x).simplify()
config_con1


# In[8]:


config_con2 = loop.dot(N.y).simplify()
config_con2


# Time differentiate the two configuation constraints. This effectively provides two motion constraints which can be used to elminate all time derivatives related to $q_2$ and $q_3$.

# In[9]:


t = me.dynamicsymbols._t


# In[10]:


config_con1_dot = config_con1.diff(t).subs({q1.diff(): u1})
config_con2_dot = config_con2.diff(t).subs({q1.diff(): u1})


# These motion constraints are linear in the $\dot{q}$'s and thus can be solved. I went ahead a plugged in $u_1$ for $\dot{q}_1$.

# In[11]:


qdots_nonmin = sm.solve([config_con1_dot, config_con2_dot], q2.diff(), q3.diff())
qdots_nonmin


# Now set the velocities and angular velocities to be only in terms of the three non-minimal q's and the single actual u.

# In[12]:


A.set_ang_vel(N, u1*N.z)
B.set_ang_vel(A, -qdots_nonmin[q2.diff()]*A.z)
C.set_ang_vel(B, -qdots_nonmin[q3.diff()]*B.z)


# In[13]:


P1.set_vel(N, 0)
P2.v2pt_theory(P1, N, A)


# In[14]:


P3.v2pt_theory(P2, N, B)


# In[15]:


P2.acc(N)


# In[16]:


P3.acc(N)


# In[17]:


P3.set_acc(N, P3.acc(N).subs(qdots_nonmin))


# In[18]:


P3.acc(N)


# # Generalized Active Forces
# 
# Apply gravity to both particles and also add a driving torque to the $l_1$ bar.

# In[19]:


RP2 = -m2*g*N.y + T/l1*A.y
RP3 = -m3*g*N.y


# In[20]:


F1 = P2.vel(N).diff(u1, N).dot(RP2) + P3.vel(N).diff(u1, N).dot(RP3)
F1


# # Generalized Inertia Forces

# In[21]:


F1star = P2.vel(N).diff(u1, N).dot(-m2*P2.acc(N)) + P3.vel(N).diff(u1, N).dot(-m3*P3.acc(N))
F1star


# # Kane's Equations
# 
# Make sure these only contain $q_1,q_2,q_3$, $u_1,\dot{u}_1$, and $T$. There should be no u's related to the non-minimal coordinates.

# In[22]:


kanes = F1 + F1star


# In[23]:


me.find_dynamicsymbols(kanes)


# In[24]:


kanes.free_symbols


# In[25]:


M = kanes.diff(u1.diff())
M


# In[26]:


G = (M*u1.diff() - kanes).simplify()


# In[27]:


me.find_dynamicsymbols(G)


# # Create a numerical function for evaluating the equations of motion

# In[28]:


from pydy.codegen.ode_function_generators import generate_ode_function


# In[29]:


eval_rhs = generate_ode_function(sm.Matrix([G]),  # right hand side of M*u1'=G
                                 (q1, q2, q3),  # all three coordinates
                                 (u1,),  # the single generalized speed
                                 (g, l1, l2, l3, m2, m3),  # constants
                                 mass_matrix=sm.Matrix([M]),  # M of M*u1'=G
                                 coordinate_derivatives=sm.Matrix([u1,  # kinematical differential equations: H of q'= H
                                                                   qdots_nonmin[q2.diff()],
                                                                   qdots_nonmin[q3.diff()]]),
                                 specifieds=(T,))  # the specified torque


# This creates a function that is useable with `scipy.integrate.odeint` or other ODE integrators.

# In[30]:


get_ipython().run_line_magic('pinfo', 'eval_rhs')


# Check to see that it works as expected:

# In[31]:


import numpy as np
rnd = np.random.random


# In[32]:


eval_rhs(rnd(4), 0.3, rnd(1), rnd(6))


# # Integrate the Equations of Motion

# In[33]:


from scipy.integrate import odeint


# Create a parameter vector:

# In[34]:


g_val = 9.81
l1_val = 0.8
l2_val = 2.0
l3_val = 1.0
l4_val = 2.0
m2_val = 1.0
m3_val = 1.0


# In[35]:


p = np.array([g_val, l1_val, l2_val, l3_val, m2_val, m3_val])


# Create an initial condition vector. Since we have configuration constraints we can only specify $q_1$ and we must calculate what $q_2$ and $q_3$ are using the configuration constraints.

# In[36]:


q1_0 = np.deg2rad(15.0)


# These two lines create a function that evaluates the configuration constraints and that will work with `scipy.optimize.fsolve` so that we can numerically solve the constraints.

# In[37]:


eval_config_con = sm.lambdify((q1, q2, q3, l1, l2, l3, l4), sm.Matrix([config_con1, config_con2]))


# In[38]:


eval_config_con_fsolve = lambda x, q1, l1, l2, l3, l4: np.squeeze(eval_config_con(q1, x[0], x[1], l1, l2, l3, l4))


# In[39]:


from scipy.optimize import fsolve


# In[40]:


q2_0, q3_0 = fsolve(eval_config_con_fsolve, np.ones(2), args=(q1_0, l1_val, l2_val, l3_val, l4_val))


# Integrate the equations with a constant $T$.

# In[41]:


times = np.linspace(0.0, 10.0, num=60*10)


# In[42]:


x0 = [q1_0, q2_0, q3_0, 0.0]
x0


# In[ ]:





# Create a function that turns on the torque for 1 second then turns it off.

# In[43]:


def step_pulse(x, t):
    if t < 2.0:
        T = 10.0
    else:
        T = 0.0
    return np.array([T])


# In[44]:


x = odeint(eval_rhs, x0, times, args=({T: step_pulse}, p))


# # Plot the results

# In[45]:


import matplotlib.pyplot as plt


# In[46]:


fig, ax = plt.subplots()
lines = ax.plot(times, np.rad2deg(x))
ax.legend([sm.latex(v, mode='inline') for v in (q1, q2, q3, u1)])
ax.set_xlabel('Time [s]')
ax.set_ylabel('[deg] or [deg/s]')


# # Animate the results with matplotlib

# In[47]:


p2_xy = np.array([l1_val*np.cos(x[:, 0]),
                  l1_val*np.sin(x[:, 0])])
p3_xy = p2_xy + np.array([l2_val*np.cos(x[:, 0] - x[:, 1]),
                          l2_val*np.sin(x[:, 0] - x[:, 1])])


# In[61]:


fig, ax = plt.subplots()
line, = ax.plot([0.0, p2_xy[0, 0], p3_xy[0, 0], l4_val],
               [0.0, p2_xy[1, 0], p3_xy[1, 0], 0.0])
title = 'Time = {:0.1f} seconds'
ax.set_title(title.format(0.0))
ax.set_ylim((-1.0, 3.0))
ax.set_xlim((-1.0, 3.0))
ax.set_aspect('equal')


# In[62]:


from matplotlib.animation import FuncAnimation

def update(i):
    xdata = [0.0, p2_xy[0, i], p3_xy[0, i], l4_val]
    ydata = [0.0, p2_xy[1, i], p3_xy[1, i], 0.0]
    line.set_data(xdata, ydata)
    ax.set_title(title.format(times[i]))
    return line,

ani = FuncAnimation(fig, update, save_count=len(times))


# In[63]:


from IPython.display import HTML
HTML(ani.to_jshtml())


# In[64]:


ani.save('four-bar-linkage.mp4', fps=60)


# In[ ]: