{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem 6 Bieber's Segway"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from numpy import zeros, cos, sin, arange, around, hstack\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib import animation\n",
    "from matplotlib.patches import Rectangle\n",
    "import numpy as np\n",
    "from scipy.interpolate import interp1d\n",
    "import sympy as smpy\n",
    "import scipy as sp\n",
    "\n",
    "# frames per second in simulation\n",
    "fps = 20\n",
    "# length of the segway arm/stick\n",
    "stick_length = 1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preamble\n",
    "This function will take care of animating the segway."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def animate_segway(t, states, controls, length):\n",
    "    #Animates the segway\n",
    "    \n",
    "    # Set up the figure, the axis, and the plot elements we want to animate\n",
    "    fig = plt.figure()\n",
    "    \n",
    "    # some config\n",
    "    segway_width = 0.4\n",
    "    segway_height = 0.2\n",
    "    \n",
    "    # x coordinate of the segway stick\n",
    "    segwayStick_x = length * np.add(states[:, 0],sin(states[:, 2]))\n",
    "    segwayStick_y = length * cos(states[:, 2])\n",
    "    \n",
    "    # set the limits\n",
    "    xmin = min(around(states[:, 0].min() - segway_width / 2.0, 1), around(segwayStick_x.min(), 1))\n",
    "    xmax = max(around(states[:, 0].max() + segway_height / 2.0, 1), around(segwayStick_y.max(), 1))\n",
    "    \n",
    "    # create the axes\n",
    "    ax = plt.axes(xlim=(xmin-.2, xmax+.2), ylim=(-length-.1, length+.1), aspect='equal')\n",
    "    \n",
    "    # display the current time\n",
    "    time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)\n",
    "    \n",
    "    # display the current control\n",
    "    control_text = ax.text(0.05, 0.8, '', transform=ax.transAxes)\n",
    "    \n",
    "    # create rectangle for the segway\n",
    "    rect = Rectangle([states[0, 0] - segway_width / 2.0, -segway_height / 2],\n",
    "        segway_width, segway_height, fill=True, color='gold', ec='blue')\n",
    "    ax.add_patch(rect)\n",
    "    \n",
    "    # blank line for the stick with o for the ends\n",
    "    stick_line, = ax.plot([], [], lw=2, marker='o', markersize=6, color='blue')\n",
    "    \n",
    "    # vector for the control (force)\n",
    "    force_vec = ax.quiver([],[],[],[],angles='xy',scale_units='xy',scale=1)\n",
    "\n",
    "    # initialization function: plot the background of each frame\n",
    "    def init():\n",
    "        time_text.set_text('')\n",
    "        control_text.set_text('')\n",
    "        rect.set_xy((0.0, 0.0))\n",
    "        stick_line.set_data([], [])\n",
    "        return time_text, rect, stick_line, control_text\n",
    "\n",
    "    # animation function: update the objects\n",
    "    def animate(i):\n",
    "        time_text.set_text('time = {:2.2f}'.format(t[i]))\n",
    "        control_text.set_text('force = {:2.3f}'.format(controls[i]))\n",
    "        rect.set_xy((states[i, 0] - segway_width / 2.0, -segway_height / 2))\n",
    "        stick_line.set_data([states[i, 0], segwayStick_x[i]], [0, segwayStick_y[i]])\n",
    "        return time_text, rect, stick_line, control_text\n",
    "\n",
    "    # call the animator function\n",
    "    anim = animation.FuncAnimation(fig, animate, frames=len(t), init_func=init,\n",
    "            interval=1000/fps, blit=False, repeat=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Dynamics: state to state\n",
    "A = np.array([[1, 0.05, -.01, 0],\n",
    "               [0, 0.22, -.17, -.01],\n",
    "               [0, 0.1, 1.14, 0.10],\n",
    "               [0, 1.66, 2.85, 1.14]]);\n",
    "# Control to state\n",
    "b = np.array([.01, .21, -.03, -0.44])\n",
    "nr_states = b.shape[0]\n",
    "\n",
    "# Initial state\n",
    "state0 = np.array([-0.3853493, 6.1032227, 0.8120005,-14])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Here is some space for you to work out the required control inputs (if any) for parts f), g) and h):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plug in your controller here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "controls = np.array([0,0,0,0]) # here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# This will add an extra couple of seconds to the simulation after the input controls with no control\n",
    "# the effect of this is just to show how the system will continue after the controller \"stops controlling\"\n",
    "controls = np.append(controls,[0, 0])\n",
    "\n",
    "# number of steps in the simulation\n",
    "nr_steps = controls.shape[0]\n",
    "\n",
    "# We now compute finer dynamics and control vectors for smoother visualization\n",
    "Afine = sp.linalg.fractional_matrix_power(A,(1/fps))\n",
    "Asum = np.eye(nr_states)\n",
    "for i in range(1, fps):\n",
    "    Asum = Asum + np.linalg.matrix_power(Afine,i)\n",
    "    \n",
    "bfine = np.linalg.inv(Asum).dot(b)\n",
    "\n",
    "# We also expand the controls in the \"intermediate steps\" (only for visualization)\n",
    "controls_final = np.outer(controls, np.ones(fps)).flatten()\n",
    "controls_final = np.append(controls_final, [0])\n",
    "\n",
    "# We compute all the states starting from x0 and using the controls\n",
    "states = np.empty([fps*(nr_steps)+1, nr_states])\n",
    "states[0,:] = state0;\n",
    "for stepId in range(1,fps*(nr_steps)+1):\n",
    "    states[stepId, :] = np.dot(Afine,states[stepId-1, :]) + controls_final[stepId-1] * bfine\n",
    "    \n",
    "# Now create the time vector for simulation\n",
    "t = np.linspace(1/fps,nr_steps,fps*(nr_steps),endpoint=True)\n",
    "t = np.append([0], t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib nbagg\n",
    "#%matplotlib qt\n",
    "animate_segway(t, states, controls_final, stick_length)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}