{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "509a68dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy as sm\n",
    "import sympy.physics.mechanics as me\n",
    "me.init_vprinting()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96569bdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "m, g, kt, kl, l = sm.symbols('m, g, k_t, k_l, l')\n",
    "q1, q2, q3 = me.dynamicsymbols('q1, q2, q3')\n",
    "u1, u2, u3 = me.dynamicsymbols('u1, u2, u3')\n",
    "\n",
    "N = me.ReferenceFrame('N')\n",
    "A = me.ReferenceFrame('A')\n",
    "B = me.ReferenceFrame('B')\n",
    "\n",
    "A.orient_axis(N, q1, N.z)\n",
    "B.orient_axis(A, q2, A.x)\n",
    "\n",
    "A.set_ang_vel(N, u1*N.z)\n",
    "B.set_ang_vel(A, u2*A.x)\n",
    "\n",
    "O = me.Point('O')\n",
    "Ao = me.Point('A_O')\n",
    "Bo = me.Point('B_O')\n",
    "\n",
    "Ao.set_pos(O, l/2*A.x)\n",
    "Bo.set_pos(O, l*A.x)\n",
    "\n",
    "O.set_vel(N, 0)\n",
    "Ao.v2pt_theory(O, N, A)\n",
    "Bo.v2pt_theory(O, N, A)\n",
    "\n",
    "Ao.vel(N), Bo.vel(N), A.ang_vel_in(N), B.ang_vel_in(N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "160fa0f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "Q = me.Point('Q')\n",
    "\n",
    "Q.set_pos(Bo, q3*B.y)\n",
    "\n",
    "Q.set_vel(B, u3*B.y)\n",
    "\n",
    "Q.v1pt_theory(Bo, N, B)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b00e62a",
   "metadata": {},
   "outputs": [],
   "source": [
    "Ao.acc(N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfe9a2ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "Bo.acc(N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "879b5f4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "Q.acc(N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a625df8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "qd_repl = {q1.diff(): u1, q2.diff(): u2, q3.diff(): u3}\n",
    "qd_repl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a360ddc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "Q.set_acc(N, Q.acc(N).xreplace(qd_repl))\n",
    "Q.acc(N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25bbea48",
   "metadata": {},
   "outputs": [],
   "source": [
    "A.ang_acc_in(N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cdb0651",
   "metadata": {},
   "outputs": [],
   "source": [
    "B.ang_acc_in(N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4065952d",
   "metadata": {},
   "outputs": [],
   "source": [
    "v_Ao_1 = Ao.vel(N).diff(u1, N)\n",
    "v_Ao_2 = Ao.vel(N).diff(u2, N)\n",
    "v_Ao_3 = Ao.vel(N).diff(u3, N)\n",
    "\n",
    "v_Ao_1, v_Ao_2, v_Ao_3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bd31eee",
   "metadata": {},
   "outputs": [],
   "source": [
    "v_Bo_1 = Bo.vel(N).diff(u1, N)\n",
    "v_Bo_2 = Bo.vel(N).diff(u2, N)\n",
    "v_Bo_3 = Bo.vel(N).diff(u3, N)\n",
    "\n",
    "v_Bo_1, v_Bo_2, v_Bo_3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4d9a507",
   "metadata": {},
   "outputs": [],
   "source": [
    "v_Q_1 = Q.vel(N).diff(u1, N)\n",
    "v_Q_2 = Q.vel(N).diff(u2, N)\n",
    "v_Q_3 = Q.vel(N).diff(u3, N)\n",
    "\n",
    "v_Q_1, v_Q_2, v_Q_3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "074b1274",
   "metadata": {},
   "outputs": [],
   "source": [
    "w_A_1 = A.ang_vel_in(N).diff(u1, N)\n",
    "w_A_2 = A.ang_vel_in(N).diff(u2, N)\n",
    "w_A_3 = A.ang_vel_in(N).diff(u3, N)\n",
    "\n",
    "w_A_1, w_A_2, w_A_3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c18f2d80",
   "metadata": {},
   "outputs": [],
   "source": [
    "w_B_1 = B.ang_vel_in(N).diff(u1, N)\n",
    "w_B_2 = B.ang_vel_in(N).diff(u2, N)\n",
    "w_B_3 = B.ang_vel_in(N).diff(u3, N)\n",
    "\n",
    "w_B_1, w_B_2, w_B_3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa515390",
   "metadata": {},
   "outputs": [],
   "source": [
    "I = m*l**2/12\n",
    "I_A_Ao = I*me.outer(A.y, A.y) + I*me.outer(A.z, A.z)\n",
    "I_B_Bo = I*me.outer(B.x, B.x) + I*me.outer(B.z, B.z)\n",
    "\n",
    "I_A_Ao, I_B_Bo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e38939d",
   "metadata": {},
   "outputs": [],
   "source": [
    "F = me.dynamicsymbols('F')\n",
    "\n",
    "R_Ao = m*g*N.x\n",
    "R_Bo = m*g*N.x - F*B.y\n",
    "R_Q = m/4*g*N.x + F*B.y\n",
    "\n",
    "R_Ao, R_Bo, R_Q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e64f2d0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "T_A = -kt*q1*N.z + kt*q2*B.x\n",
    "T_B = -kt*q2*B.x\n",
    "\n",
    "T_A, T_B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cebcbd08",
   "metadata": {},
   "outputs": [],
   "source": [
    "Rs_Ao = -m*Ao.acc(N)\n",
    "Rs_Bo = -m*Bo.acc(N)\n",
    "Rs_Q = -m/4*Q.acc(N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cc96540",
   "metadata": {},
   "outputs": [],
   "source": [
    "Ts_A = -(me.dot(A.ang_acc_in(N), I_A_Ao) + \n",
    "         me.dot(me.cross(A.ang_vel_in(N), I_A_Ao), A.ang_vel_in(N)))\n",
    "Ts_A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6014023",
   "metadata": {},
   "outputs": [],
   "source": [
    "Ts_B = -(me.dot(B.ang_acc_in(N), I_B_Bo) + \n",
    "         me.dot(me.cross(B.ang_vel_in(N), I_B_Bo), B.ang_vel_in(N)))\n",
    "Ts_B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72010e03",
   "metadata": {},
   "outputs": [],
   "source": [
    "me.find_dynamicsymbols(Ts_A, reference_frame=N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26b704cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "me.find_dynamicsymbols(Ts_B, reference_frame=N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dfb5828",
   "metadata": {},
   "outputs": [],
   "source": [
    "Fr = sm.Matrix([\n",
    "    v_Ao_1.dot(R_Ao) + v_Bo_1.dot(R_Bo) + v_Q_1.dot(R_Q) +\n",
    "    w_A_1.dot(T_A) + w_B_1.dot(T_B),\n",
    "    v_Ao_2.dot(R_Ao) + v_Bo_2.dot(R_Bo) + v_Q_2.dot(R_Q) +\n",
    "    w_A_2.dot(T_A) + w_B_2.dot(T_B),\n",
    "    v_Ao_3.dot(R_Ao) + v_Bo_3.dot(R_Bo) + v_Q_3.dot(R_Q) +\n",
    "    w_A_3.dot(T_A) + w_B_3.dot(T_B),\n",
    "])\n",
    "Fr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "495260dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "Frs = sm.Matrix([\n",
    "    v_Ao_1.dot(Rs_Ao) + v_Bo_1.dot(Rs_Bo) + v_Q_1.dot(Rs_Q) +\n",
    "    w_A_1.dot(Ts_A) + w_B_1.dot(Ts_B),\n",
    "    v_Ao_2.dot(Rs_Ao) + v_Bo_2.dot(Rs_Bo) + v_Q_2.dot(Rs_Q) +\n",
    "    w_A_2.dot(Ts_A) + w_B_2.dot(Ts_B),\n",
    "    v_Ao_3.dot(Rs_Ao) + v_Bo_3.dot(Rs_Bo) + v_Q_3.dot(Rs_Q) +\n",
    "    w_A_3.dot(Ts_A) + w_B_3.dot(Ts_B),\n",
    "])\n",
    "Frs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41ec1afe",
   "metadata": {},
   "outputs": [],
   "source": [
    "kanes_eqs = Fr + Frs\n",
    "kanes_eqs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06d3e714",
   "metadata": {},
   "outputs": [],
   "source": [
    "u = sm.Matrix([u1, u2, u3])\n",
    "t = me.dynamicsymbols._t\n",
    "ud = u.diff(t)\n",
    "ud"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae102805",
   "metadata": {},
   "outputs": [],
   "source": [
    "Md = Frs.jacobian(ud)\n",
    "Md"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5b94664",
   "metadata": {},
   "outputs": [],
   "source": [
    "ud_zero_repl = {u1.diff(): 0, u2.diff(): 0, u3.diff(): 0}\n",
    "ud_zero_repl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e21c23e",
   "metadata": {},
   "outputs": [],
   "source": [
    "gd = Frs.xreplace(ud_zero_repl) + Fr\n",
    "gd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4ab1ef4",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = sm.cse(Md.LUsolve(-gd))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a839a67c",
   "metadata": {},
   "outputs": [],
   "source": [
    "res[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8de7bb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "res[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf9be36a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}