def test_one_dof():
    # This is for a 1 dof spring-mass-damper case.
    # It is described in more detail in the Kane docstring.
    q, u = dynamicsymbols('q u')
    qd, ud = dynamicsymbols('q u', 1)
    m, c, k = symbols('m c k')
    N = ReferenceFrame('N')
    P = Point('P')
    P.set_vel(N, u * N.x)

    kd = [qd - u]
    FL = [(P, (-k * q - c * u) * N.x)]
    pa = Particle()
    pa.mass = m
    pa.point = P
    BL = [pa]

    KM = Kane(N)
    KM.coords([q])
    KM.speeds([u])
    KM.kindiffeq(kd)
    KM.kanes_equations(FL, BL)
    MM = KM.mass_matrix
    forcing = KM.forcing
    rhs = MM.inv() * forcing
    assert expand(rhs[0]) == expand(-(q * k + u * c) / m)
    assert KM.linearize() == (Matrix([[0, 1], [k, c]]), Matrix([]))

def test_pend():
    q, u = dynamicsymbols('q u')
    qd, ud = dynamicsymbols('q u', 1)
    m, l, g = symbols('m l g')
    N = ReferenceFrame('N')
    P = Point('P')
    P.set_vel(N, -l * u * sin(q) * N.x + l * u * cos(q) * N.y)
    kd = [qd - u]

    FL = [(P, m * g * N.x)]
    pa = Particle()
    pa.mass = m
    pa.point = P
    BL = [pa]

    KM = Kane(N)
    KM.coords([q])
    KM.speeds([u])
    KM.kindiffeq(kd)
    KM.kanes_equations(FL, BL)
    MM = KM.mass_matrix
    forcing = KM.forcing
    rhs = MM.inv() * forcing
    rhs.simplify()
    assert expand(rhs[0]) == expand(-g / l * sin(q))

def test_particle():
    m = Symbol('m')
    P = Point('P')
    p = Particle()
    assert p.mass == None
    assert p.point == None
    # Test the mass setter
    p.mass = m
    assert p.mass == m
    # Test the point setter
    p.point = P
    assert p.point == P

def test_particle():
    m, m2 = symbols("m m2")
    P = Point("P")
    P2 = Point("P2")
    p = Particle("pa", P, m)
    assert p.mass == m
    assert p.point == P
    # Test the mass setter
    p.mass = m2
    assert p.mass == m2
    # Test the point setter
    p.point = P2
    assert p.point == P2

def test_dub_pen():

    # The system considered is the double pendulum. Like in the
    # test of the simple pendulum above, we begin by creating the generalized
    # coordinates and the simple generalized speeds and accelerations which
    # will be used later. Following this we create frames and points necessary
    # for the kinematics. The procedure isn't explicitly explained as this is
    # similar to the simple  pendulum. Also this is documented on the pydy.org
    # website.
    q1, q2 = dynamicsymbols('q1 q2')
    q1d, q2d = dynamicsymbols('q1 q2', 1)
    q1dd, q2dd = dynamicsymbols('q1 q2', 2)
    u1, u2 = dynamicsymbols('u1 u2')
    u1d, u2d = dynamicsymbols('u1 u2', 1)
    l, m, g = symbols('l m g')

    N = ReferenceFrame('N')
    A = N.orientnew('A', 'Axis', [q1, N.z])
    B = N.orientnew('B', 'Axis', [q2, N.z])

    A.set_ang_vel(N, q1d * A.z)
    B.set_ang_vel(N, q2d * A.z)

    O = Point('O')
    P = O.locatenew('P', l * A.x)
    R = P.locatenew('R', l * B.x)

    O.set_vel(N, 0)
    P.v2pt_theory(O, N, A)
    R.v2pt_theory(P, N, B)

    ParP = Particle('ParP', P, m)
    ParR = Particle('ParR', R, m)

    ParP.potential_energy = - m * g * l * cos(q1)
    ParR.potential_energy = - m * g * l * cos(q1) - m * g * l * cos(q2)
    L = Lagrangian(N, ParP, ParR)
    lm = LagrangesMethod(L, [q1, q2], bodies=[ParP, ParR])
    lm.form_lagranges_equations()

    assert simplify(l*m*(2*g*sin(q1) + l*sin(q1)*sin(q2)*q2dd
        + l*sin(q1)*cos(q2)*q2d**2 - l*sin(q2)*cos(q1)*q2d**2
        + l*cos(q1)*cos(q2)*q2dd + 2*l*q1dd) - lm.eom[0]) == 0
    assert simplify(l*m*(g*sin(q2) + l*sin(q1)*sin(q2)*q1dd
        - l*sin(q1)*cos(q2)*q1d**2 + l*sin(q2)*cos(q1)*q1d**2
        + l*cos(q1)*cos(q2)*q1dd + l*q2dd) - lm.eom[1]) == 0
    assert lm.bodies == [ParP, ParR]

def test_potential_energy():
    m, M, l1, g, h, H = symbols('m M l1 g h H')
    omega = dynamicsymbols('omega')
    N = ReferenceFrame('N')
    O = Point('O')
    O.set_vel(N, 0 * N.x)
    Ac = O.locatenew('Ac', l1 * N.x)
    P = Ac.locatenew('P', l1 * N.x)
    a = ReferenceFrame('a')
    a.set_ang_vel(N, omega * N.z)
    Ac.v2pt_theory(O, N, a)
    P.v2pt_theory(O, N, a)
    Pa = Particle('Pa', P, m)
    I = outer(N.z, N.z)
    A = RigidBody('A', Ac, a, M, (I, Ac))
    Pa.set_potential_energy(m * g * h)
    A.set_potential_energy(M * g * H)
    assert potential_energy(A, Pa) == m * g * h + M * g * H

def test_potential_energy():
    m, M, l1, g, h, H = symbols("m M l1 g h H")
    omega = dynamicsymbols("omega")
    N = ReferenceFrame("N")
    O = Point("O")
    O.set_vel(N, 0 * N.x)
    Ac = O.locatenew("Ac", l1 * N.x)
    P = Ac.locatenew("P", l1 * N.x)
    a = ReferenceFrame("a")
    a.set_ang_vel(N, omega * N.z)
    Ac.v2pt_theory(O, N, a)
    P.v2pt_theory(O, N, a)
    Pa = Particle("Pa", P, m)
    I = outer(N.z, N.z)
    A = RigidBody("A", Ac, a, M, (I, Ac))
    Pa.potential_energy = m * g * h
    A.potential_energy = M * g * H
    assert potential_energy(A, Pa) == m * g * h + M * g * H

    def __init__(self, name, masscenter=None, mass=None, frame=None,
                 central_inertia=None):

        self.name = name
        self.loads = []

        if frame is None:
            frame = ReferenceFrame(name + '_frame')

        if masscenter is None:
            masscenter = Point(name + '_masscenter')

        if central_inertia is None and mass is None:
            ixx = Symbol(name + '_ixx')
            iyy = Symbol(name + '_iyy')
            izz = Symbol(name + '_izz')
            izx = Symbol(name + '_izx')
            ixy = Symbol(name + '_ixy')
            iyz = Symbol(name + '_iyz')
            _inertia = (inertia(frame, ixx, iyy, izz, ixy, iyz, izx),
                        masscenter)
        else:
            _inertia = (central_inertia, masscenter)

        if mass is None:
            _mass = Symbol(name + '_mass')
        else:
            _mass = mass

        masscenter.set_vel(frame, 0)

        # If user passes masscenter and mass then a particle is created
        # otherwise a rigidbody. As a result a body may or may not have inertia.
        if central_inertia is None and mass is not None:
            self.frame = frame
            self.masscenter = masscenter
            Particle.__init__(self, name, masscenter, _mass)
        else:
            RigidBody.__init__(self, name, masscenter, frame, _mass, _inertia)

def test_simp_pen():
    # This tests that the equations generated by LagrangesMethod are identical
    # to those obtained by hand calculations. The system under consideration is
    # the simple pendulum.
    # We begin by creating the generalized coordinates as per the requirements
    # of LagrangesMethod. Also we created the associate symbols
    # that characterize the system: 'm' is the mass of the bob, l is the length
    # of the massless rigid rod connecting the bob to a point O fixed in the
    # inertial frame.
    q, u = dynamicsymbols('q u')
    qd, ud = dynamicsymbols('q u ', 1)
    l, m, g = symbols('l m g')

    # We then create the inertial frame and a frame attached to the massless
    # string following which we define the inertial angular velocity of the
    # string.
    N = ReferenceFrame('N')
    A = N.orientnew('A', 'Axis', [q, N.z])
    A.set_ang_vel(N, qd * N.z)

    # Next, we create the point O and fix it in the inertial frame. We then
    # locate the point P to which the bob is attached. Its corresponding
    # velocity is then determined by the 'two point formula'.
    O = Point('O')
    O.set_vel(N, 0)
    P = O.locatenew('P', l * A.x)
    P.v2pt_theory(O, N, A)

    # The 'Particle' which represents the bob is then created and its
    # Lagrangian generated.
    Pa = Particle('Pa', P, m)
    Pa.set_potential_energy(- m * g * l * cos(q))
    L = Lagrangian(N, Pa)

    # The 'LagrangesMethod' class is invoked to obtain equations of motion.
    lm = LagrangesMethod(L, [q])
    lm.form_lagranges_equations()
    RHS = lm.rhs()
    assert RHS[1] == -g*sin(q)/l

def test_lagrange_2forces():
    ### Equations for two damped springs in serie with two forces

    ### generalized coordinates
    qs = q1, q2 = dynamicsymbols('q1, q2')
    ### generalized speeds
    qds = q1d, q2d = dynamicsymbols('q1, q2', 1)

    ### Mass, spring strength, friction coefficient
    m, k, nu = symbols('m, k, nu')

    N = ReferenceFrame('N')
    O = Point('O')

    ### Two points
    P1 = O.locatenew('P1', q1 * N.x)
    P1.set_vel(N, q1d * N.x)
    P2 = O.locatenew('P1', q2 * N.x)
    P2.set_vel(N, q2d * N.x)

    pP1 = Particle('pP1', P1, m)
    pP1.potential_energy = k * q1**2 / 2

    pP2 = Particle('pP2', P2, m)
    pP2.potential_energy = k * (q1 - q2)**2 / 2

    #### Friction forces
    forcelist = [(P1, - nu * q1d * N.x),
                 (P2, - nu * q2d * N.x)]
    lag = Lagrangian(N, pP1, pP2)

    l_method = LagrangesMethod(lag, (q1, q2), forcelist=forcelist, frame=N)
    l_method.form_lagranges_equations()

    eq1 = l_method.eom[0]
    assert eq1.diff(q1d) == nu
    eq2 = l_method.eom[1]
    assert eq2.diff(q2d) == nu

def test_two_dof():
    # This is for a 2 d.o.f., 2 particle spring-mass-damper.
    # The first coordinate is the displacement of the first particle, and the
    # second is the relative displacement between the first and second
    # particles. Speeds are defined as the time derivatives of the particles.
    q1, q2, u1, u2 = dynamicsymbols('q1 q2 u1 u2')
    q1d, q2d, u1d, u2d = dynamicsymbols('q1 q2 u1 u2', 1)
    m, c1, c2, k1, k2 = symbols('m c1 c2 k1 k2')
    N = ReferenceFrame('N')
    P1 = Point('P1')
    P2 = Point('P2')
    P1.set_vel(N, u1 * N.x)
    P2.set_vel(N, (u1 + u2) * N.x)
    kd = [q1d - u1, q2d - u2]

    # Now we create the list of forces, then assign properties to each
    # particle, then create a list of all particles.
    FL = [(P1, (-k1 * q1 - c1 * u1 + k2 * q2 + c2 * u2) * N.x), (P2, (-k2 *
        q2 - c2 * u2) * N.x)]
    pa1 = Particle()
    pa2 = Particle()
    pa1.mass = m
    pa2.mass = m
    pa1.point = P1
    pa2.point = P2
    BL = [pa1, pa2]

    # Finally we create the Kane object, specify the inertial frame, pass
    # relevant information, and form Fr & Fr*. Then we calculate the mass
    # matrix and forcing terms, and finally solve for the udots.
    KM = Kane(N)
    KM.coords([q1, q2])
    KM.speeds([u1, u2])
    KM.kindiffeq(kd)
    KM.kanes_equations(FL, BL)
    MM = KM.mass_matrix
    forcing = KM.forcing
    rhs = MM.inv() * forcing
    assert expand(rhs[0]) == expand((-k1 * q1 - c1 * u1 + k2 * q2 + c2 * u2)/m)
    assert expand(rhs[1]) == expand((k1 * q1 + c1 * u1 - 2 * k2 * q2 - 2 *
                                    c2 * u2) / m)

def test_particle():
    m, m2, v1, v2, v3, r, g, h = symbols('m m2 v1 v2 v3 r g h')
    P = Point('P')
    P2 = Point('P2')
    p = Particle('pa', P, m)
    assert p.mass == m
    assert p.point == P
    # Test the mass setter
    p.mass = m2
    assert p.mass == m2
    # Test the point setter
    p.point = P2
    assert p.point == P2
    # Test the linear momentum function
    N = ReferenceFrame('N')
    O = Point('O')
    P2.set_pos(O, r * N.y)
    P2.set_vel(N, v1 * N.x)
    assert p.linear_momentum(N) == m2 * v1 * N.x
    assert p.angular_momentum(O, N) == -m2 * r *v1 * N.z
    P2.set_vel(N, v2 * N.y)
    assert p.linear_momentum(N) == m2 * v2 * N.y
    assert p.angular_momentum(O, N) == 0
    P2.set_vel(N, v3 * N.z)
    assert p.linear_momentum(N) == m2 * v3 * N.z
    assert p.angular_momentum(O, N) == m2 * r * v3 * N.x
    P2.set_vel(N, v1 * N.x + v2 * N.y + v3 * N.z)
    assert p.linear_momentum(N) == m2 * (v1 * N.x + v2 * N.y + v3 * N.z)
    assert p.angular_momentum(O, N) == m2 * r * (v3 * N.x - v1 * N.z)
    p.set_potential_energy(m * g * h)
    assert p.potential_energy == m * g * h
    # TODO make the result not be system-dependent
    assert p.kinetic_energy(
        N) in [m2*(v1**2 + v2**2 + v3**2)/2,
        m2 * v1**2 / 2 + m2 * v2**2 / 2 + m2 * v3**2 / 2]

m, g, l = symbols('m g l')
N = ReferenceFrame('N')

# part a
r1 = s*N.x
r2 = (s + l*cos(theta))*N.x + l*sin(theta)*N.y

O = Point('O')
p1 = O.locatenew('p1', r1)
p2 = O.locatenew('p2', r2)

O.set_vel(N, 0)
p1.set_vel(N, p1.pos_from(O).dt(N))
p2.set_vel(N, p2.pos_from(O).dt(N))

P1 = Particle('P1', p1, 2*m)
P2 = Particle('P2', p2, m)

P1.set_potential_energy(0)
P2.set_potential_energy(P2.mass * g * (p2.pos_from(O) & N.y))

L1 = Lagrangian(N, P1, P2)
print('{} = {}'.format('L1', msprint(L1)))

lm1 = LagrangesMethod(L1, [s, theta])
lm1.form_lagranges_equations()
rhs = lm1.rhs()
t = symbols('t')
print('{} = {}'.format(msprint(sd.diff(t)), msprint(rhs[2].simplify())))
print('{} = {}\n'.format(msprint(thetad.diff(t)), msprint(rhs[3].simplify())))

def test_particle():
    m, m2, v1, v2, v3, r = symbols('m m2 v1 v2 v3 r')
    P = Point('P')
    P2 = Point('P2')
    p = Particle('pa', P, m)
    assert p.mass == m
    assert p.point == P
    # Test the mass setter
    p.mass = m2
    assert p.mass == m2
    # Test the point setter
    p.point = P2
    assert p.point == P2
    # Test the linear momentum function
    N = ReferenceFrame('N')
    O = Point('O')
    P2.set_pos(O, r * N.y)
    P2.set_vel(N, v1 * N.x)
    assert p.linearmomentum(N) == m2 * v1 * N.x
    assert p.angularmomentum(O, N) == -m2 * r *v1 * N.z
    P2.set_vel(N, v2 * N.y)
    assert p.linearmomentum(N) == m2 * v2 * N.y
    assert p.angularmomentum(O, N) == 0
    P2.set_vel(N, v3 * N.z)
    assert p.linearmomentum(N) == m2 * v3 * N.z
    assert p.angularmomentum(O, N) == m2 * r * v3 * N.x
    P2.set_vel(N, v1 * N.x + v2 * N.y + v3 * N.z)
    assert p.linearmomentum(N) == m2 * (v1 * N.x + v2 * N.y + v3 * N.z)
    assert p.angularmomentum(O, N) == m2 * r * (v3 * N.x - v1 * N.z)

one_frame.ang_vel_in(inertial_frame)
two_frame.set_ang_vel(one_frame, omega2*inertial_frame.z)
two_frame.ang_vel_in(inertial_frame)

#Sets up the linear velocities of the points on the linkages
#one.set_vel(inertial_frame, 0)
two.v2pt_theory(one, inertial_frame, one_frame)
two.vel(inertial_frame)
three.v2pt_theory(two, inertial_frame, two_frame)
three.vel(inertial_frame)

#Sets up the masses of the linkages
one_mass, two_mass = symbols('m_1, m_2')

#Defines the linkages as particles
twoP = Particle('twoP', two, one_mass)
threeP = Particle('threeP', three, two_mass)

#Sets up gravity information and assigns gravity to act on mass centers
g = symbols('g')
two_grav_force_vector = -1*one_mass*g*inertial_frame.y
two_grav_force = (two, two_grav_force_vector)
three_grav_force_vector = -1*two_mass*g*inertial_frame.y
three_grav_force = (three, three_grav_force_vector)

#Sets up joint torques
one_torque, two_torque = dynamicsymbols('T_1, T_2')
one_torque_vector = one_torque*inertial_frame.z - two_torque*inertial_frame.z
one_link_torque = (one_frame, one_torque_vector)

two_torque_vector = two_torque*inertial_frame.z

c_frame.ang_vel_in(inertial_frame)

#Sets up the linear velocities of the points on the linkages
A.set_vel(inertial_frame, 0)
B.v2pt_theory(A, inertial_frame, a_frame)
B.vel(inertial_frame)
C.v2pt_theory(B, inertial_frame, b_frame)
C.vel(inertial_frame)
D.v2pt_theory(C, inertial_frame, c_frame)
D.vel(inertial_frame)

#Sets up the masses of the linkages
a_mass, b_mass, c_mass = symbols('m_a, m_b, m_c')

#Defines the linkages as particles
bP = Particle('bP', B, a_mass)
cP = Particle('cP', C, b_mass)
dP = Particle('dP', D, c_mass)

#Sets up gravity information and assigns gravity to act on mass centers
g = symbols('g')
b_grav_force_vector = -1*a_mass*g*inertial_frame.y
b_grav_force = (B, b_grav_force_vector)
c_grav_force_vector = -1*b_mass*g*inertial_frame.y
c_grav_force = (C, c_grav_force_vector)
d_grav_force_vector = -1*c_mass*g*inertial_frame.y
d_grav_force = (D, d_grav_force_vector)

#Sets up joint torques
a_torque, b_torque, c_torque = dynamicsymbols('T_a, T_b, T_c')
a_torque_vector = a_torque*inertial_frame.z - b_torque*inertial_frame.z