{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook provides examples to go along with the [textbook](http://underactuated.csail.mit.edu/contact.html).  I recommend having both windows open, side-by-side!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import mpld3\n",
    "import numpy as np\n",
    "from IPython.display import display\n",
    "from pydrake.all import DirectCollocation, MathematicalProgram, Solve\n",
    "from pydrake.examples.rimless_wheel import RimlessWheel\n",
    "\n",
    "from underactuated import FindResource, running_as_notebook\n",
    "\n",
    "if running_as_notebook:\n",
    "    mpld3.enable_notebook()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Trajectory optimization for the Rimless Wheel\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rimless_wheel_limit_cycle():\n",
    "    plant = RimlessWheel()\n",
    "    context = plant.CreateDefaultContext()\n",
    "\n",
    "    params = context.get_numeric_parameter(0)\n",
    "    slope = params.slope()\n",
    "    alpha = np.pi / params.number_of_spokes()\n",
    "\n",
    "    dircol = DirectCollocation(plant,\n",
    "                            context,\n",
    "                            num_time_samples=15,\n",
    "                            minimum_timestep=0.01,\n",
    "                            maximum_timestep=0.1,\n",
    "                            assume_non_continuous_states_are_fixed=True)\n",
    "    prog = dircol.prog()\n",
    "\n",
    "    dircol.AddEqualTimeIntervalsConstraints()\n",
    "\n",
    "    dircol.AddConstraintToAllKnotPoints(dircol.state()[0] >= slope - alpha)\n",
    "    dircol.AddConstraintToAllKnotPoints(dircol.state()[0] <= slope + alpha)\n",
    "\n",
    "    prog.AddConstraint(dircol.initial_state()[0] == slope - alpha)\n",
    "    prog.AddConstraint(dircol.final_state()[0] == slope + alpha)\n",
    "\n",
    "    prog.AddConstraint(dircol.initial_state()[1] == dircol.final_state()[1]\n",
    "                       * np.cos(2 * alpha))\n",
    "\n",
    "    result = Solve(prog)\n",
    "    assert result.is_success()\n",
    "\n",
    "    x_trajectory = dircol.ReconstructStateTrajectory(result)\n",
    "\n",
    "    x_knots = np.hstack([\n",
    "        x_trajectory.value(t) for t in np.linspace(x_trajectory.start_time(),\n",
    "                                                x_trajectory.end_time(), 100)\n",
    "    ])\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.plot(x_knots[0, :], x_knots[1, :])\n",
    "    ax.plot([x_knots[0, 0], x_knots[0, -1]], [x_knots[1, 0], x_knots[1, -1]], \"--\")\n",
    "\n",
    "    # Plot the energy contours.\n",
    "    nq = 151\n",
    "    nqd = 151\n",
    "    mgl = params.mass() * params.gravity() * params.length()\n",
    "    q = np.linspace(-0.5, 0.5, nq)\n",
    "    qd = np.linspace(-.5, 2, nqd)\n",
    "    Q, QD = np.meshgrid(q, qd)\n",
    "    Energy = .5 * params.mass() * params.length()**2 * QD**2 + mgl * np.cos(Q)\n",
    "    ax.contour(Q,\n",
    "            QD,\n",
    "            Energy,\n",
    "            alpha=0.5,\n",
    "            linestyles=\"dashed\",\n",
    "            colors=\"black\",\n",
    "            linewidths=0.5)\n",
    "\n",
    "    ax.set_xlabel(\"theta\")\n",
    "    ax.set_ylabel(\"thetadot\")\n",
    "    ax.axis([-0.5, 0.5, 0, 2])\n",
    "    ax.set_title(\"Limit Cycle of the Rimless Wheel (w/ contours of \"\n",
    "                \"constant energy)\")\n",
    "    display(mpld3.display())\n",
    "\n",
    "rimless_wheel_limit_cycle()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A simple basketball \"bounce pass\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "e = .8    # coefficient of restitution (0 ≤ e ≤ 1)\n",
    "g = 9.8  # gravitational constant (m/s^2)\n",
    "x0 = 0\n",
    "xf = 4\n",
    "z0 = 1\n",
    "zf = 1\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10,4))\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('z')\n",
    "ax.set_xlim(x0-.1, xf+.1)\n",
    "\n",
    "def bounce_pass(number_of_bounces=0, zdot0_is_positive=True, debug=False):\n",
    "    \n",
    "    duration = 2+number_of_bounces\n",
    "    xdot = (xf-x0)/duration\n",
    "    prog = MathematicalProgram()\n",
    "    h = prog.NewContinuousVariables(number_of_bounces+1, 'h')\n",
    "    prog.AddConstraint(np.sum(h) == duration)\n",
    "    prog.AddBoundingBoxConstraint(0, np.inf, h)\n",
    "    \n",
    "    # zdot at the start of each segment (after collision).\n",
    "    zdot = prog.NewContinuousVariables(number_of_bounces+1, 'zd')\n",
    "\n",
    "    # Initial velocity constraint (since we *might* have two solutions for each number of bounces)\n",
    "    if (zdot0_is_positive):\n",
    "        prog.AddBoundingBoxConstraint(0, np.inf, zdot[0])\n",
    "    else:\n",
    "        prog.AddBoundingBoxConstraint(-np.inf, 0, zdot[0])\n",
    "    \n",
    "    # Add dynamics constraints for each segment that ends with a bounce\n",
    "    for i in range(number_of_bounces): \n",
    "        # z must be zero at the end of this segment:\n",
    "        z = zdot[i]*h[i] - .5*g*h[i]*h[i] + (z0 if i==0 else 0)\n",
    "        prog.AddConstraint(z == 0)\n",
    "        prog.AddConstraint(zdot[i+1] == -e*(zdot[i] - h[i]*g))\n",
    "        \n",
    "    # Now the final segment\n",
    "    z = zdot[-1]*h[-1] - .5*g*h[-1]*h[-1] + (z0 if number_of_bounces==0 else 0)\n",
    "    prog.AddConstraint(z == zf)\n",
    "    \n",
    "    result = Solve(prog)\n",
    "    if not result.is_success():\n",
    "        if debug:\n",
    "            infeasible = result.GetInfeasibleConstraints(prog)\n",
    "            print(\"Infeasible constraints:\")\n",
    "            for i in range(len(infeasible)):\n",
    "                print(infeasible[i])\n",
    "        # return without plotting\n",
    "        return\n",
    "\n",
    "    # plot the resulting trajectory\n",
    "    ax.set_prop_cycle(plt.rcParams['axes.prop_cycle']) # reset color cycle\n",
    "    relative_time = np.linspace(0, 1, 10)\n",
    "    x_start = x0\n",
    "    for i in range(number_of_bounces+1):\n",
    "        t = result.GetSolution(h[i])*relative_time\n",
    "        x = x_start + t*xdot\n",
    "        z = (z0 if i==0 else 0) + result.GetSolution(zdot[i])*t - .5*g*t*t\n",
    "        ax.plot(x, z, '.-')\n",
    "        x_start = x[-1]\n",
    "\n",
    "#bounce_pass(number_of_bounces=2, zdot0_is_positive=False, debug=True)\n",
    "for number_of_bounces in range(0,4):\n",
    "    bounce_pass(number_of_bounces=number_of_bounces, zdot0_is_positive=True)\n",
    "    bounce_pass(number_of_bounces=number_of_bounces, zdot0_is_positive=False)\n",
    "\n",
    "display(mpld3.display())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The basketball \"trick shot\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trick_shot():\n",
    "    m = 1.\n",
    "    r = 0.1\n",
    "    g = 9.8  # gravitational constant (m/s^2)\n",
    "    e = .8    # coefficient of restitution (0 ≤ e ≤ 1)\n",
    "    x0 = -.25\n",
    "    xf = -1\n",
    "    z0 = 1\n",
    "    zf = 3\n",
    "\n",
    "    N = 3\n",
    "    prog = MathematicalProgram()\n",
    "    h = prog.NewContinuousVariables(N)\n",
    "    # Make sure time intervals are positive:\n",
    "    prog.AddBoundingBoxConstraint(0.01, 10, h)\n",
    "        \n",
    "    # zdot at the start of each segment (after collision).\n",
    "    xdot = prog.NewContinuousVariables(N)\n",
    "    zdot = prog.NewContinuousVariables(N)\n",
    "    thetadot = prog.NewContinuousVariables(N)\n",
    "\n",
    "    # First segment -- from the launch to the wall.  \n",
    "    # Implement the guard (x==0).\n",
    "    prog.AddConstraint(x0 + h[0]*xdot[0] == 0).evaluator().set_description('x_wall')\n",
    "    z = z0 + h[0]*zdot[0] - 0.5*g*h[0]*h[0]\n",
    "    zdot_pre = zdot[0] - g*h[0]\n",
    "    # integration + collision dynamics\n",
    "    prog.AddConstraint(xdot[1] == -e * xdot[0]).evaluator().set_description('xdot_wall')\n",
    "    prog.AddConstraint(zdot[1] == 3.*zdot_pre/5. - 2.*r*thetadot[0]/5.).evaluator().set_description('zdot_wall')\n",
    "    prog.AddConstraint(thetadot[1] == -3*zdot_pre/(5.*r) + 2.*thetadot[0]/5.).evaluator().set_description('thetadot_wall')\n",
    "    # Help it find the solution we want\n",
    "    prog.AddConstraint(zdot[0] >= 0)\n",
    "    prog.AddConstraint(thetadot[0] <= 0)\n",
    "\n",
    "    # Second segment -- from wall to the floor\n",
    "    z = z + zdot[1]*h[1] - 0.5*g*h[1]*h[1]\n",
    "    # Implement the guard (z==0).\n",
    "    prog.AddConstraint(z == 0).evaluator().set_description('z_floor')\n",
    "    x = xdot[1]*h[1]\n",
    "    zdot_pre = zdot[1] - g*h[1]\n",
    "    # integration + collision dynamics\n",
    "    prog.AddConstraint(xdot[2] == 3.*xdot[1]/5. - 2.*r*thetadot[1]/5.).evaluator().set_description('xdot_floor')\n",
    "    prog.AddConstraint(zdot[2] == -e*zdot_pre).evaluator().set_description('zdot_floor')\n",
    "    prog.AddConstraint(thetadot[2] == -3*xdot[1]/(5.*r) + 2.*thetadot[1]/5.).evaluator().set_description('thetadot_floor')\n",
    "    # Make sure the ball doesn't travel too far away\n",
    "    prog.AddConstraint(x >= -3.)\n",
    "\n",
    "    # Final segment -- from the floor to the hoop\n",
    "    x = x + xdot[2]*h[2]\n",
    "    z = zdot[2]*h[2] - 0.5*g*h[2]*h[2]\n",
    "    # make sure we end at the hoop.\n",
    "    prog.AddConstraint(x == xf).evaluator().set_description('x_hoop')\n",
    "    prog.AddConstraint(z == zf).evaluator().set_description('z_hoop')\n",
    "    zdot_final = zdot[2] - g*h[2]\n",
    "    # make sure we're approaching from the correct direction.\n",
    "    prog.AddConstraint(xdot[2] >= .1).evaluator().set_description('xdot_hoop')\n",
    "    prog.AddConstraint(zdot_final <= -2).evaluator().set_description('zdot_hoop')\n",
    "\n",
    "    result = Solve(prog)\n",
    "    if not result.is_success():\n",
    "        infeasible = result.GetInfeasibleConstraints(prog)\n",
    "        print(\"Infeasible constraints:\")\n",
    "        for i in range(len(infeasible)):\n",
    "            print(infeasible[i])\n",
    "\n",
    "        # plot the resulting trajectory\n",
    "    fig, ax = plt.subplots(figsize=(10,4))\n",
    "    ax.set_xlabel('x')\n",
    "    ax.set_ylabel('z')\n",
    "    #ax.set_xlim(x0-.1, xf+.1)\n",
    "    ax.set_prop_cycle(plt.rcParams['axes.prop_cycle']) # reset color cycle\n",
    "    relative_time = np.linspace(0, 1, 10)\n",
    "    x_start = x0\n",
    "    z_start = z0\n",
    "    for i in range(N):\n",
    "        t = result.GetSolution(h[i])*relative_time\n",
    "        x = x_start + result.GetSolution(xdot[i])*t\n",
    "        z = z_start + result.GetSolution(zdot[i])*t - .5*g*t*t\n",
    "        ax.plot(x, z, '.-')\n",
    "        x_start = x[-1]\n",
    "        z_start = z[-1]\n",
    "        \n",
    "    print('Rotation rate during each segment:')\n",
    "    print(result.GetSolution(thetadot))\n",
    "    display(mpld3.display())\n",
    "\n",
    "trick_shot()"
   ]
  },
  {
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}