from sympy import init_session
init_session(use_latex=True, quiet=True)

M_b = Symbol('M_b')  # Masa de la pala
g = Symbol('g')  # Gravedad
e = Symbol('e')  # Excentricidad de batimiento

from sympy.physics.mechanics import ReferenceFrame, Point, dynamicsymbols

# Definimos nuestra propia clase para que los versores sean IJK
class IJKReferenceFrame(ReferenceFrame):
    def __init__(self, name):
        super().__init__(name, latexs=['\mathbf{%s}_{%s}' % (idx, name) for idx in ("i", "j", "k")])

A = IJKReferenceFrame("A")

# Definimos la velocidad de giro del rotor principal $\Omega$:
Omega = Symbol('Omega')
Omega

A1 = IJKReferenceFrame("A1")
psi = dynamicsymbols('psi')
A1.orient(A, 'Axis', [psi, A.z])
A1.dcm(A)  # T_{A1A}

T_a, T_k, T_h = dynamicsymbols('T_a, T_k, T_h')
T_a

t = symbols('t')
omega1, omega2, omega3 = dynamicsymbols('omega1, omega2, omega3')
theta1, theta2, theta3 = dynamicsymbols('theta1, theta2, theta3')
eqs = [omega1 - theta1.diff(t), omega2 - theta2.diff(t), omega3 - theta3.diff(t)]
eqs

psi.diff() * A.z

A1.ang_vel_in(A)

A2 = IJKReferenceFrame("A2")
beta = dynamicsymbols('beta')
A2.orient(A1, 'Axis', [beta, -A1.y])
A2.dcm(A1)  # T_{A2A1}

A2.ang_vel_in(A1)

A2.dcm(A)  # T_{A2A}

A2.ang_vel_in(A)

A3 = IJKReferenceFrame("A3")
zeta = dynamicsymbols('zeta')
A3.orient(A2, 'Axis', [zeta, A2.z])
A3.dcm(A2)  # T_{A3A2}

A3.dcm(A1)  # T_{A3A1}

B = IJKReferenceFrame("B")
theta = dynamicsymbols('theta')
B.orient(A3, 'Axis', [theta, A3.x])
B.dcm(A3)  # T_{BA3}

B.dcm(A1)  # T_{BA1}

B.dcm(A)  # T_{BA}

O = Point("O")
O.set_vel(A, 0)

E = O.locatenew('E', e * A1.x)
E.pos_from(O)

x_GB = Symbol('x_GB')
GB = E.locatenew('GB', x_GB * B.x)
GB.pos_from(E)

omega_B = B.ang_vel_in(A).subs(psi.diff(), Omega)
omega_B

omega_B

omega_B.to_matrix(B)

from sympy.physics.mechanics import inertia

I_theta, I_beta, I_zeta = symbols('I_theta, I_beta, I_zeta')
I_B = inertia(B, I_theta, I_beta, I_zeta)
I_B

I_B.to_matrix(B)

h_B = I_B.dot(omega_B).subs(psi.diff(), Omega)
h_B.to_matrix(B)

dh_B_dt = h_B.dt(A)
dh_B_dt

E.set_vel(A1, 0)
v_E = E.v1pt_theory(O, A, A1)
v_E

a_E = v_E.dt(A).subs({psi.diff(): Omega, psi.diff().diff(): 0})
a_E

a_E.to_matrix(B).simplify()

EGB = GB.pos_from(E)
EGB.cross(M_b * a_E).to_matrix(B).simplify()

EGB.cross(M_b * (-g * A1.z)).to_matrix(B).simplify()

omega_B.express(B).subs({theta: 0, zeta: 0})

_.doit()

h_B.express(B).subs({theta: 0, zeta: 0}).doit()

dh_B_dt

(dh_B_dt.express(B)
    .subs({
           theta: 0,
           zeta: 0,
           psi.diff(): Omega,
           psi.diff().diff(): 0,
           I_zeta: I_theta + I_beta})
    .doit()
    .to_matrix(B)
    .simplify()
    .expand(trig=True))

EGB.cross(M_b * a_E).express(B).subs({theta: 0, zeta: 0}).doit()

EGB.cross(M_b * (-g * A1.z)).express(B).subs({theta: 0, zeta: 0}).doit()

MrxB, MrzB = symbols('M^r_xB, M^r_zB')
k_beta = Symbol('k_beta')
Mr = MrxB * B.x + k_beta * beta * B.y + MrzB * B.z
Mr

lhs = Mr + EGB.cross(M_b * (-g * A1.z)).express(B).subs({theta: 0, zeta: 0}).doit()
rhs = (dh_B_dt.express(B).subs({theta: 0, zeta: 0, psi.diff(): Omega, psi.diff().diff(): 0, I_zeta: I_theta + I_beta}).doit()
       + EGB.cross(M_b * a_E).express(B).subs({theta: 0, zeta: 0}).doit())
lhs.simplify()
rhs.simplify()
eqxB = Eq((lhs & B.x) - (rhs & B.x))
eqxB

eqyB = Eq(expand_trig(collect((lhs & B.y) - (rhs & B.y), I_beta)))
eqyB

eqzB = Eq((lhs & B.z) - (rhs & B.z))
eqzB

beta_ = Dummy('beta')
eqyB_lin = Eq((eqyB.lhs
               .subs(beta, beta_)
               .series(beta_, n=2)
               .removeO()
               .subs(beta_, beta)
               / I_beta)
              .expand()
              .collect(beta))
eqyB_lin

assert eqyB_lin.lhs.is_Add
# Filtramos término gravitatorio
args = [add for add in eqyB_lin.lhs.args if not add.has(g)]
# Reconstruimos la ecuación
eqyB_lin = Eq(eqyB_lin.lhs.func(*args))
eqyB_lin

eqyB_lin = Eq((eqyB_lin.lhs / Omega ** 2)
              .expand()
              .collect(beta)
              * Omega ** 2)

assert eqyB_lin.lhs.is_Mul
for factor in eqyB_lin.lhs.args:
    if factor.is_Add:
        adds = factor.args
    else:
        mul = factor

eqyB_lin = Eq(mul * adds[0] + mul * adds[1])
eqyB_lin

Omega_beta = Symbol('Omega_beta')
eqyB_lin = Eq(eqyB_lin.lhs.subs(
    Omega * sqrt(1 + M_b * e * x_GB / I_beta + k_beta / (I_beta * Omega ** 2)),
    Omega_beta))
eqyB_lin

try:
    Eq(Omega * A.z, 3 * A.z)
except Exception as e:
    print(e)

from sympy.abc import a
Eq(a, 1)

Eq(a, 1) + 1

from sympy.physics.mechanics import RigidBody

Pala = RigidBody('Pala', GB, B, M_b, [I_B, E])

Pala.central_inertia

E.set_vel(A, E.v2pt_theory(O, A, A1))
Pala.angular_momentum(E, A)  # ¿Qué estoy intentando hallar?