{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Problem 4 Bieber's Segway"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the following block of code first to get all the dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# %load gauss_elim.py\n",
    "from gauss_elim import gauss_elim"
   ]
  },
  {
   "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 scipy as sp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dynamics"
   ]
  },
  {
   "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])\n",
    "\n",
    "# Final (terminal state)\n",
    "stateFinal = np.array([0, 0, 0, 0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part (f), (g), (h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# You may use gauss_elim to help you find the row reduced echelon form."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part (i)"
   ]
  },
  {
   "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": [
    "# frames per second in simulation\n",
    "fps = 20\n",
    "# length of the segway arm/stick\n",
    "stick_length = 1.\n",
    "\n",
    "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": [
    "### 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": "markdown",
   "metadata": {},
   "source": [
    "### Simulation"
   ]
  },
  {
   "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": "markdown",
   "metadata": {},
   "source": [
    "### Visualization"
   ]
  },
  {
   "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": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [default]",
   "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.4.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}