#!/usr/bin/env python

from __future__ import print_function

import collections
import matplotlib.animation as animation
import matplotlib.pyplot as plt
import numpy as np
import time
from pid import PIDController

m = 0.1895
n = 0.16
c = 0.14
g = 9.81
l = 0.2221
# Update frequency, frames per second, decides accuracy of simulation
update_freq = 30.0
# Number of seconds to remember back in second plot
time_length = 30
num_samples = int(update_freq * time_length)

class Joint(object):
    """ Simulated joint """
    def __init__(self):
        self.theta = np.random.uniform(0, np.pi)
        self.vel = 0.0
        self.current = 0.0
        self.contr = PIDController()

    def update(self, desired, dt):
        """ Update simulation """
        nu = n * self.current - c * self.vel
        mgl = m * g * l * (np.cos(self.theta) * -1.0)
        ml = 2.0 * m * (l ** 2)
        acc = (nu + mgl) / ml
        self.theta += self.vel * dt + 0.5 * acc * (dt ** 2)
        self.vel += acc * dt
        if 0.0 >= self.theta or self.theta >= np.pi:
            self.theta = max(0.0, min(np.pi, self.theta))
            self.vel = 0.0
        try:
            self.current = self.contr(self.theta - np.pi / 2.0,
                    desired - np.pi / 2.0, self.vel, dt)
            self.current = max(-4.1, min(4.1, self.current))
        except NotImplementedError:
            # Ignore if not implemented
            pass

class JointSim(object):
    """ Joint simulation class """
    def __init__(self):
        self.theta_d = np.pi / 2.0
        self.joint = Joint()
        self.last_update = time.time()
        self.paused = False

    def prepare(self):
        """ Prepare figure for animation """
        fig = plt.figure()
        # Create subplot for pendulum
        self.ax = fig.add_subplot(2, 1, 1)
        self.ax.axis([-1, 1, 0, 1])
        self.actual, = self.ax.plot(np.array([0.0, np.cos(self.joint.theta)]),
                np.array([0.0, np.sin(self.joint.theta)]), linewidth = 2)
        self.desired, = self.ax.plot(np.array([0.0, np.cos(self.theta_d)]),
                np.array([0.0, np.sin(self.theta_d)]), '--')
        fig.canvas.mpl_connect('button_release_event', self._onclick)
        fig.canvas.mpl_connect('key_release_event', self._keypress)
        # Create subplot for error over time
        ax2 = fig.add_subplot(2, 1, 2)
        ax2.axis([0, num_samples, -np.pi / 2.0, np.pi / 2.0])
        ax2.set_ylabel('Angle error')
        ax2.set_xlabel('Time')
        ax2.set_xticks([])
        ax2.plot([0, num_samples], [0, 0], 'g--')
        self.error = collections.deque(maxlen = num_samples)
        self.error_plt, = ax2.plot(self.error, 'b-')
        # `blit` can be turned to `True` for better performance, but because
        # of a race condition in matplotlib it is not enabled by default
        self.ani = animation.FuncAnimation(fig, self._update, init_func=self._update,
                blit=False, interval=int(1000.0 / update_freq))

    def _update(self, i=None):
        """ Internal method to update animation """
        now = time.time()
        # Don't allow to large updates steps, this prevents obvious jitter and
        # non-linear effects in acceleration should not present them self.
        # The problem is that this can cause "slow-motion" movement which will
        # look unrealistic.
        if not self.paused:
            delta = min(now - self.last_update, 0.04)
            self.joint.update(self.theta_d, delta)
            self.desired.set_xdata(np.array([0.0, np.cos(self.theta_d)]))
            self.desired.set_ydata(np.array([0.0, np.sin(self.theta_d)]))
            self.actual.set_xdata(np.array([0.0, np.cos(self.joint.theta)]))
            self.actual.set_ydata(np.array([0.0, np.sin(self.joint.theta)]))
            # Add new error data
            self.error.append(self.joint.theta - self.theta_d)
            self.error_plt.set_data(range(len(self.error)), self.error)
        self.last_update = now
        return self.actual, self.desired, self.error_plt,

    def _onclick(self, event):
        """ Handle click events on canvas """
        if event.xdata is not None and event.ydata is not None and event.button == 1:
            if event.inaxes == self.ax:
                self.theta_d = np.arctan2(event.ydata, event.xdata)
            else:
                current_time = (event.xdata / num_samples) * time_length
                index = int((event.xdata / num_samples) * num_samples)
                value = None
                if index < len(self.error):
                    value = self.error[int((event.xdata / num_samples) * num_samples)]
                print(current_time, value)

    def _keypress(self, event):
        """ Handle keyboard release events """
        if event.key is not None and event.key == ' ':
            self.paused = not self.paused

if __name__ == '__main__':
    # Create simulation
    sim = JointSim()
    # Prepare drawing and animation
    sim.prepare()
    # Tell Matplotlib to start drawing
    plt.show()