def test_linear_momentum():
    N = ReferenceFrame("N")
    Ac = Point("Ac")
    Ac.set_vel(N, 25 * N.y)
    I = outer(N.x, N.x)
    A = RigidBody("A", Ac, N, 20, (I, Ac))
    P = Point("P")
    Pa = Particle("Pa", P, 1)
    Pa.point.set_vel(N, 10 * N.x)
    assert linear_momentum(N, A, Pa) == 10 * N.x + 500 * N.y

def test_rigidbody2():
    M, v, r, omega = dynamicsymbols('M v r omega')
    N = ReferenceFrame('N')
    b = ReferenceFrame('b')
    b.set_ang_vel(N, omega * b.x)
    P = Point('P')
    I = outer (b.x, b.x)
    Inertia_tuple = (I, P)
    B = RigidBody('B', P, b, M, Inertia_tuple)
    P.set_vel(N, v * b.x)
    assert B.angularmomentum(P, N) == omega * b.x
    O = Point('O')
    O.set_vel(N, v * b.x)
    P.set_pos(O, r * b.y)
    assert B.angularmomentum(O, N) == omega * b.x - M*v*r*b.z

def test_kinetic_energy():
    m, M, l1 = symbols('m M l1')
    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))
    assert 0 == kinetic_energy(N, Pa, A) - (M*l1**2*omega**2/2
            + 2*l1**2*m*omega**2 + omega**2/2)

def test_gravity():
    N = ReferenceFrame('N')
    m, M, g = symbols('m M g')
    F1, F2 = dynamicsymbols('F1 F2')
    po = Point('po')
    pa = Particle('pa', po, m)
    A = ReferenceFrame('A')
    P = Point('P')
    I = outer(A.x, A.x)
    B = RigidBody('B', P, A, M, (I, P))
    forceList = [(po, F1), (P, F2)]
    forceList.extend(gravity(g*N.y, pa, B))
    l = [(po, F1), (P, F2), (po, g*m*N.y), (P, g*M*N.y)]

    for i in range(len(l)):
        for j in range(len(l[i])):
            assert forceList[i][j] == l[i][j]

def test_angular_momentum_and_linear_momentum():
    m, M, l1 = symbols('m M l1')
    q1d = dynamicsymbols('q1d')
    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, q1d * 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))
    assert linear_momentum(N, A, Pa) == 2 * m * q1d* l1 * N.y + M * l1 * q1d * N.y
    assert angular_momentum(O, N, A, Pa) == 4 * m * q1d * l1**2 * N.z + q1d * N.z

def test_kinetic_energy():
    m, M, l1 = symbols("m M l1")
    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))
    assert 0 == kinetic_energy(N, Pa, A) - (
        M * l1 ** 2 * omega ** 2 / 2 + 2 * l1 ** 2 * m * omega ** 2 + omega ** 2 / 2
    )

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_rigidbody2():
    M, v, r, omega, g, h = dynamicsymbols('M v r omega g h')
    N = ReferenceFrame('N')
    b = ReferenceFrame('b')
    b.set_ang_vel(N, omega * b.x)
    P = Point('P')
    I = outer(b.x, b.x)
    Inertia_tuple = (I, P)
    B = RigidBody('B', P, b, M, Inertia_tuple)
    P.set_vel(N, v * b.x)
    assert B.angular_momentum(P, N) == omega * b.x
    O = Point('O')
    O.set_vel(N, v * b.x)
    P.set_pos(O, r * b.y)
    assert B.angular_momentum(O, N) == omega * b.x - M*v*r*b.z
    B.set_potential_energy(M * g * h)
    assert B.potential_energy == M * g * h
    assert B.kinetic_energy(N) == (omega**2 + M * v**2) / 2

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 test_center_of_mass():
    a = ReferenceFrame('a')
    m = symbols('m', real=True)
    p1 = Particle('p1', Point('p1_pt'), S(1))
    p2 = Particle('p2', Point('p2_pt'), S(2))
    p3 = Particle('p3', Point('p3_pt'), S(3))
    p4 = Particle('p4', Point('p4_pt'), m)
    b_f = ReferenceFrame('b_f')
    b_cm = Point('b_cm')
    mb = symbols('mb')
    b = RigidBody('b', b_cm, b_f, mb, (outer(b_f.x, b_f.x), b_cm))
    p2.point.set_pos(p1.point, a.x)
    p3.point.set_pos(p1.point, a.x + a.y)
    p4.point.set_pos(p1.point, a.y)
    b.masscenter.set_pos(p1.point, a.y + a.z)
    point_o=Point('o')
    point_o.set_pos(p1.point, center_of_mass(p1.point, p1, p2, p3, p4, b))
    expr = 5/(m + mb + 6)*a.x + (m + mb + 3)/(m + mb + 6)*a.y + mb/(m + mb + 6)*a.z
    assert point_o.pos_from(p1.point)-expr == 0

def test_rigidbody3():
    q1, q2, q3, q4 = dynamicsymbols('q1:5')
    p1, p2, p3 = symbols('p1:4')
    m = symbols('m')

    A = ReferenceFrame('A')
    B = A.orientnew('B', 'axis', [q1, A.x])
    O = Point('O')
    O.set_vel(A, q2*A.x + q3*A.y + q4*A.z)
    P = O.locatenew('P', p1*B.x + p2*B.y + p3*B.z)
    I = outer(B.x, B.x)

    rb1 = RigidBody('rb1', P, B, m, (I, P))
    # I_S/O = I_S/S* + I_S*/O
    rb2 = RigidBody('rb2', P, B, m,
                    (I + inertia_of_point_mass(m, P.pos_from(O), B), O))

    assert rb1.central_inertia == rb2.central_inertia
    assert rb1.angular_momentum(O, A) == rb2.angular_momentum(O, A)

def test_angular_momentum_and_linear_momentum():
    """A rod with length 2l, centroidal inertia I, and mass M along with a
    particle of mass m fixed to the end of the rod rotate with an angular rate
    of omega about point O which is fixed to the non-particle end of the rod.
    The rod's reference frame is A and the inertial frame is N."""
    m, M, l, I = symbols('m, M, l, I')
    omega = dynamicsymbols('omega')
    N = ReferenceFrame('N')
    a = ReferenceFrame('a')
    O = Point('O')
    Ac = O.locatenew('Ac', l * N.x)
    P = Ac.locatenew('P', l * N.x)
    O.set_vel(N, 0 * N.x)
    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)
    A = RigidBody('A', Ac, a, M, (I * outer(N.z, N.z), Ac))
    expected = 2 * m * omega * l * N.y + M * l * omega * N.y
    assert linear_momentum(N, A, Pa) == expected
    expected = (I + M * l**2 + 4 * m * l**2) * omega * N.z
    assert angular_momentum(O, N, A, Pa) == expected

def test_angular_momentum_and_linear_momentum():
    """A rod with length 2l, centroidal inertia I, and mass M along with a
    particle of mass m fixed to the end of the rod rotate with an angular rate
    of omega about point O which is fixed to the non-particle end of the rod.
    The rod's reference frame is A and the inertial frame is N."""
    m, M, l, I = symbols("m, M, l, I")
    omega = dynamicsymbols("omega")
    N = ReferenceFrame("N")
    a = ReferenceFrame("a")
    O = Point("O")
    Ac = O.locatenew("Ac", l * N.x)
    P = Ac.locatenew("P", l * N.x)
    O.set_vel(N, 0 * N.x)
    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)
    A = RigidBody("A", Ac, a, M, (I * outer(N.z, N.z), Ac))
    expected = 2 * m * omega * l * N.y + M * l * omega * N.y
    assert (linear_momentum(N, A, Pa) - expected) == Vector(0)
    expected = (I + M * l ** 2 + 4 * m * l ** 2) * omega * N.z
    assert (angular_momentum(O, N, A, Pa) - expected).simplify() == Vector(0)

import sympy.physics.mechanics as me
import sympy as sm
import math as m
import numpy as np

frame_a = me.ReferenceFrame('a')
frame_b = me.ReferenceFrame('b')
q1, q2, q3 = me.dynamicsymbols('q1 q2 q3')
frame_b.orient(frame_a, 'Axis', [q3, frame_a.x])
dcm = frame_a.dcm(frame_b)
m = dcm*3-frame_a.dcm(frame_b)
r = me.dynamicsymbols('r')
circle_area = sm.pi*r**2
u, a = me.dynamicsymbols('u a')
x, y = me.dynamicsymbols('x y')
s = u*me.dynamicsymbols._t-1/2*a*me.dynamicsymbols._t**2
expr1 = 2*a*0.5-1.25+0.25
expr2 = -1*x**2+y**2+0.25*(x+y)**2
expr3 = 0.5*10**(-10)
dyadic=me.outer(frame_a.x, frame_a.x)+me.outer(frame_a.y, frame_a.y)+me.outer(frame_a.z, frame_a.z)

point_o = me.Point('o')
point_p = me.Point('p')
point_o.set_pos(point_p, c1*frame_a.x)
v = (v).express(frame_n)
point_o.set_pos(point_p, (point_o.pos_from(point_p)).express(frame_n))
frame_a.set_ang_vel(frame_n, c3*frame_a.z)
print(frame_n.ang_vel_in(frame_a))
point_p.v2pt_theory(point_o,frame_n,frame_a)
particle_p1 = me.Particle('p1', me.Point('p1_pt'), sm.Symbol('m'))
particle_p2 = me.Particle('p2', me.Point('p2_pt'), sm.Symbol('m'))
particle_p2.point.v2pt_theory(particle_p1.point,frame_n,frame_a)
point_p.a2pt_theory(particle_p1.point,frame_n,frame_a)
body_b1_cm = me.Point('b1_cm')
body_b1_cm.set_vel(frame_n, 0)
body_b1_f = me.ReferenceFrame('b1_f')
body_b1 = me.RigidBody('b1', body_b1_cm, body_b1_f, sm.symbols('m'), (me.outer(body_b1_f.x,body_b1_f.x),body_b1_cm))
body_b2_cm = me.Point('b2_cm')
body_b2_cm.set_vel(frame_n, 0)
body_b2_f = me.ReferenceFrame('b2_f')
body_b2 = me.RigidBody('b2', body_b2_cm, body_b2_f, sm.symbols('m'), (me.outer(body_b2_f.x,body_b2_f.x),body_b2_cm))
g = sm.symbols('g', real=True)
force_p1 = particle_p1.mass*(g*frame_n.x)
force_p2 = particle_p2.mass*(g*frame_n.x)
force_b1 = body_b1.mass*(g*frame_n.x)
force_b2 = body_b2.mass*(g*frame_n.x)
z = me.dynamicsymbols('z')
v=x*frame_a.x+y*frame_a.z
point_o.set_pos(point_p, x*frame_a.x+y*frame_a.y)
v = (v).subs({x:2*z, y:z})
point_o.set_pos(point_p, (point_o.pos_from(point_p)).subs({x:2*z, y:z}))
force_o = -1*(x*y*frame_a.x)

import sympy.physics.mechanics as me
import sympy as sm
import math as m
import numpy as np

g, lb, w, h=sm.symbols('g lb w h', real=True)
theta, phi, omega, alpha = me.dynamicsymbols('theta phi omega alpha')
thetad, phid, omegad, alphad = me.dynamicsymbols('theta phi omega alpha', 1)
thetad2, phid2 = me.dynamicsymbols('theta phi', 2)
frame_n=me.ReferenceFrame('n')
body_a_cm=me.Point('a_cm')
body_a_cm.set_vel(frame_n, 0)
body_a_f=me.ReferenceFrame('a_f')
body_a=me.RigidBody('a', body_a_cm, body_a_f, sm.symbols('m'), (me.outer(body_a_f.x,body_a_f.x),body_a_cm))
body_b_cm=me.Point('b_cm')
body_b_cm.set_vel(frame_n, 0)
body_b_f=me.ReferenceFrame('b_f')
body_b=me.RigidBody('b', body_b_cm, body_b_f, sm.symbols('m'), (me.outer(body_b_f.x,body_b_f.x),body_b_cm))
body_a_f.orient(frame_n, 'Axis', [theta, frame_n.y])
body_b_f.orient(body_a_f, 'Axis', [phi, body_a_f.z])
point_o=me.Point('o')
la = (lb-h/2)/2
body_a_cm.set_pos(point_o, la*body_a_f.z)
body_b_cm.set_pos(point_o, lb*body_a_f.z)
body_a_f.set_ang_vel(frame_n, omega*frame_n.y)
body_b_f.set_ang_vel(body_a_f, alpha*body_a_f.z)
point_o.set_vel(frame_n, 0)
body_a_cm.v2pt_theory(point_o,frame_n,body_a_f)
body_b_cm.v2pt_theory(point_o,frame_n,body_a_f)
ma=sm.symbols('ma')
body_a.mass = ma

frame_n = me.ReferenceFrame('n')
x1, x2, x3 = me.dynamicsymbols('x1 x2 x3')
l = sm.symbols('l', real=True)
v1=x1*frame_a.x+x2*frame_a.y+x3*frame_a.z
v2=x1*frame_b.x+x2*frame_b.y+x3*frame_b.z
v3=x1*frame_n.x+x2*frame_n.y+x3*frame_n.z
v=v1+v2+v3
point_c = me.Point('c')
point_d = me.Point('d')
point_po1 = me.Point('po1')
point_po2 = me.Point('po2')
point_po3 = me.Point('po3')
particle_l = me.Particle('l', me.Point('l_pt'), sm.Symbol('m'))
particle_p1 = me.Particle('p1', me.Point('p1_pt'), sm.Symbol('m'))
particle_p2 = me.Particle('p2', me.Point('p2_pt'), sm.Symbol('m'))
particle_p3 = me.Particle('p3', me.Point('p3_pt'), sm.Symbol('m'))
body_s_cm = me.Point('s_cm')
body_s_cm.set_vel(frame_n, 0)
body_s_f = me.ReferenceFrame('s_f')
body_s = me.RigidBody('s', body_s_cm, body_s_f, sm.symbols('m'), (me.outer(body_s_f.x,body_s_f.x),body_s_cm))
body_r1_cm = me.Point('r1_cm')
body_r1_cm.set_vel(frame_n, 0)
body_r1_f = me.ReferenceFrame('r1_f')
body_r1 = me.RigidBody('r1', body_r1_cm, body_r1_f, sm.symbols('m'), (me.outer(body_r1_f.x,body_r1_f.x),body_r1_cm))
body_r2_cm = me.Point('r2_cm')
body_r2_cm.set_vel(frame_n, 0)
body_r2_f = me.ReferenceFrame('r2_f')
body_r2 = me.RigidBody('r2', body_r2_cm, body_r2_f, sm.symbols('m'), (me.outer(body_r2_f.x,body_r2_f.x),body_r2_cm))
v4=x1*body_s_f.x+x2*body_s_f.y+x3*body_s_f.z
body_s_cm.set_pos(point_c, l*frame_n.x)

import sympy.physics.mechanics as me
import sympy as sm
import math as m
import numpy as np

frame_a=me.ReferenceFrame('a')
c1, c2, c3=sm.symbols('c1 c2 c3', real=True)
a=me.inertia(frame_a, 1, 1, 1)
particle_p1=me.Particle('p1', me.Point('p1_pt'), sm.Symbol('m'))
particle_p2=me.Particle('p2', me.Point('p2_pt'), sm.Symbol('m'))
body_r_cm=me.Point('r_cm')
body_r_f=me.ReferenceFrame('r_f')
body_r=me.RigidBody('r', body_r_cm, body_r_f, sm.symbols('m'), (me.outer(body_r_f.x,body_r_f.x),body_r_cm))
r_a = sm.Matrix([1,1,1,1,1,0,0,0,1]).reshape(3, 3)
point_o=me.Point('o')
m1=sm.symbols('m1')
particle_p1.mass = m1
m2=sm.symbols('m2')
particle_p2.mass = m2
mr=sm.symbols('mr')
body_r.mass = mr
i1 = sm.symbols('i1')
i2 = sm.symbols('i2')
i3 = sm.symbols('i3')
body_r.inertia = (me.inertia(body_r_f, i1, i2, i3, 0, 0, 0), body_r_cm)
point_o.set_pos(particle_p1.point, c1*frame_a.x)
point_o.set_pos(particle_p2.point, c2*frame_a.y)
point_o.set_pos(body_r_cm, c3*frame_a.z)
a=me.inertia_of_point_mass(particle_p1.mass, particle_p1.point.pos_from(point_o), frame_a)
a=me.inertia_of_point_mass(particle_p2.mass, particle_p2.point.pos_from(point_o), frame_a)
a=body_r.inertia[0] + me.inertia_of_point_mass(body_r.mass, body_r.masscenter.pos_from(point_o), frame_a)