{
  "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import threading\n",
    "import time\n",
    "import sys\n",
    "\n",
    "from ipywidgets import interact\n",
    "import ipywidgets as widgets\n",
    "import numpy as np\n",
    "import pinocchio as pin\n",
    "\n",
    "sys.path.append(\"..\")\n",
    "import talos_conf as conf\n",
    "import vizutils\n",
    "from tsid_biped import TsidBiped"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tsid = TsidBiped(conf)\n",
    "tsid.q0[2] = 1.02127\n",
    "\n",
    "com_0 = tsid.robot.com(tsid.formulation.data())\n",
    "H_rf_0 = tsid.robot.framePosition(\n",
    "    tsid.formulation.data(), tsid.model.getFrameId(conf.rf_frame_name)\n",
    ")\n",
    "H_lf_0 = tsid.robot.framePosition(\n",
    "    tsid.formulation.data(), tsid.model.getFrameId(conf.lf_frame_name)\n",
    ")\n",
    "\n",
    "vizutils.addViewerSphere(\n",
    "    tsid.viz, \"world/com\", conf.SPHERE_RADIUS, conf.COM_SPHERE_COLOR\n",
    ")\n",
    "vizutils.addViewerSphere(\n",
    "    tsid.viz, \"world/com_ref\", conf.REF_SPHERE_RADIUS, conf.COM_REF_SPHERE_COLOR\n",
    ")\n",
    "vizutils.addViewerSphere(tsid.viz, \"world/rf\", conf.SPHERE_RADIUS, conf.RF_SPHERE_COLOR)\n",
    "vizutils.addViewerSphere(\n",
    "    tsid.viz, \"world/rf_ref\", conf.REF_SPHERE_RADIUS, conf.RF_REF_SPHERE_COLOR\n",
    ")\n",
    "vizutils.addViewerSphere(tsid.viz, \"world/lf\", conf.SPHERE_RADIUS, conf.LF_SPHERE_COLOR)\n",
    "vizutils.addViewerSphere(\n",
    "    tsid.viz, \"world/lf_ref\", conf.REF_SPHERE_RADIUS, conf.LF_REF_SPHERE_COLOR\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "push_robot_active, push_robot_com_vel, com_vel_entry = False, 3 * [0.0], None\n",
    "run, push_robot_active = threading.Event(), threading.Event()\n",
    "run.set()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_simu():\n",
    "    i, t = 0, 0.0\n",
    "    q, v = tsid.q, tsid.v\n",
    "    time_avg = 0.0\n",
    "    while run.is_set():\n",
    "        time_start = time.time()\n",
    "\n",
    "        tsid.comTask.setReference(tsid.trajCom.computeNext())\n",
    "        tsid.postureTask.setReference(tsid.trajPosture.computeNext())\n",
    "        tsid.rightFootTask.setReference(tsid.trajRF.computeNext())\n",
    "        tsid.leftFootTask.setReference(tsid.trajLF.computeNext())\n",
    "\n",
    "        HQPData = tsid.formulation.computeProblemData(t, q, v)\n",
    "\n",
    "        sol = tsid.solver.solve(HQPData)\n",
    "        if sol.status != 0:\n",
    "            print(\"QP problem could not be solved! Error code:\", sol.status)\n",
    "            break\n",
    "\n",
    "        # tau = tsid.formulation.getActuatorForces(sol)\n",
    "        dv = tsid.formulation.getAccelerations(sol)\n",
    "        q, v = tsid.integrate_dv(q, v, dv, conf.dt)\n",
    "        i, t = i + 1, t + conf.dt\n",
    "\n",
    "        if push_robot_active.is_set():\n",
    "            push_robot_active.clear()\n",
    "            data = tsid.formulation.data()\n",
    "            if tsid.contact_LF_active:\n",
    "                J_LF = tsid.contactLF.computeMotionTask(0.0, q, v, data).matrix\n",
    "            else:\n",
    "                J_LF = np.zeros((0, tsid.model.nv))\n",
    "            if tsid.contact_RF_active:\n",
    "                J_RF = tsid.contactRF.computeMotionTask(0.0, q, v, data).matrix\n",
    "            else:\n",
    "                J_RF = np.zeros((0, tsid.model.nv))\n",
    "            J = np.vstack((J_LF, J_RF))\n",
    "            J_com = tsid.comTask.compute(t, q, v, data).matrix\n",
    "            A = np.vstack((J_com, J))\n",
    "            b = np.concatenate((np.array(push_robot_com_vel), np.zeros(J.shape[0])))\n",
    "            v = np.linalg.lstsq(A, b, rcond=-1)[0]\n",
    "\n",
    "        if i % conf.DISPLAY_N == 0:\n",
    "            tsid.viz.display(q)\n",
    "            x_com = tsid.robot.com(tsid.formulation.data())\n",
    "            x_com_ref = tsid.trajCom.getSample(t).pos()\n",
    "            H_lf = tsid.robot.framePosition(tsid.formulation.data(), tsid.LF)\n",
    "            H_rf = tsid.robot.framePosition(tsid.formulation.data(), tsid.RF)\n",
    "            x_lf_ref = tsid.trajLF.getSample(t).pos()[:3]\n",
    "            x_rf_ref = tsid.trajRF.getSample(t).pos()[:3]\n",
    "            vizutils.applyViewerConfiguration(\n",
    "                tsid.viz, \"world/com\", [*x_com.tolist(), 0, 0, 0, 1.0]\n",
    "            )\n",
    "            vizutils.applyViewerConfiguration(\n",
    "                tsid.viz, \"world/com_ref\", [*x_com_ref.tolist(), 0, 0, 0, 1.0]\n",
    "            )\n",
    "            vizutils.applyViewerConfiguration(\n",
    "                tsid.viz, \"world/rf\", pin.SE3ToXYZQUATtuple(H_rf)\n",
    "            )\n",
    "            vizutils.applyViewerConfiguration(\n",
    "                tsid.viz, \"world/lf\", pin.SE3ToXYZQUATtuple(H_lf)\n",
    "            )\n",
    "            vizutils.applyViewerConfiguration(\n",
    "                tsid.viz, \"world/rf_ref\", [*x_rf_ref.tolist(), 0, 0, 0, 1.0]\n",
    "            )\n",
    "            vizutils.applyViewerConfiguration(\n",
    "                tsid.viz, \"world/lf_ref\", [*x_lf_ref.tolist(), 0, 0, 0, 1.0]\n",
    "            )\n",
    "\n",
    "        time_spent = time.time() - time_start\n",
    "        time_avg = (i * time_avg + time_spent) / (i + 1)\n",
    "\n",
    "        if time_avg < 0.9 * conf.dt:\n",
    "            time.sleep(conf.dt - time_avg)\n",
    "    print(\"end of simulation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = widgets.Output()\n",
    "\n",
    "th_simu = threading.Thread(target=run_simu)\n",
    "th_simu.start()\n",
    "\n",
    "\n",
    "@interact(com_x=(-10.0, 10.0), com_y=(-15.0, 15.0), com_z=(-40.0, 40.0))\n",
    "def sliders_com(com_x, com_y, com_z):\n",
    "    tsid.trajCom.setReference(\n",
    "        com_0 + np.array([1e-2 * com_x, 1e-2 * com_y, 1e-2 * com_z]).T\n",
    "    )\n",
    "\n",
    "\n",
    "@interact(rf_x=(-30.0, 30.0), rf_y=(-30.0, 30.0), rf_z=(-30.0, 30.0))\n",
    "def sliders_rf(rf_x, rf_y, rf_z):\n",
    "    H_rf_ref = H_rf_0.copy()\n",
    "    H_rf_ref.translation += np.array([1e-2 * rf_x, 1e-2 * rf_y, 1e-2 * rf_z]).T\n",
    "    tsid.trajRF.setReference(H_rf_ref)\n",
    "\n",
    "\n",
    "@interact(lf_x=(-30.0, 30.0), lf_y=(-30.0, 30.0), lf_z=(-30.0, 30.0))\n",
    "def sliders_lf(lf_x, lf_y, lf_z):\n",
    "    H_lf_ref = H_lf_0.copy()\n",
    "    H_lf_ref.translation += +np.array([1e-2 * lf_x, 1e-2 * lf_y, 1e-2 * lf_z]).T\n",
    "    tsid.trajLF.setReference(H_lf_ref)\n",
    "\n",
    "\n",
    "def buttons(b):\n",
    "    if b.description == \"Stop\":\n",
    "        run.clear()\n",
    "        th_simu.join()\n",
    "    elif b.description == \"Switch Right Foot\":\n",
    "        if tsid.contact_RF_active:\n",
    "            tsid.remove_contact_RF()\n",
    "        else:\n",
    "            tsid.add_contact_RF()\n",
    "    elif b.description == \"Switch Left Foot\":\n",
    "        if tsid.contact_LF_active:\n",
    "            tsid.remove_contact_LF()\n",
    "        else:\n",
    "            tsid.add_contact_LF()\n",
    "\n",
    "\n",
    "for action in [\"Stop\", \"Switch Right Foot\", \"Switch Left Foot\"]:\n",
    "    button = widgets.Button(description=action)\n",
    "    button.on_click(buttons)\n",
    "    output.append_display_data(button)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hasattr(tsid.viz.viewer, \"jupyter_cell\") and tsid.viz.viewer.jupyter_cell()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
