import sys
from random import random
import numpy as np
import pylab
import time
import scipy.integrate

"""
  This script compares different methods to solve an ODE.

  The solved system consists of:
  y_dot = f(y, t, p)
  y0

  where
   y: the state vector of the system
   y0: the initial state of the system
   y_dot: the derivatives of the state vector
   f: the function that returns the derivatives.
   t: the time instant
   p: the parameters of the system
"""

def func(y, t, parameters):
   """
    The example system is a bouncing point mass on a hunt-cossley surface.
    State vector:
    y[0] = height
    y[1] = vertical momentum

    y_dot[0] = velocity
    y_dot[1] = force
   """
   # Step 1: extract the states into variables:
   h = y[0]
   p = y[1]
   mass = parameters[0]
   Kp_floor = parameters[1]
   Kd_floor = parameters[2]

   # Step 2: Calculate model:
   v = p / mass

   # Ground reaction force:
   if h > 0:
      F = 0.0
   else:
      F = -h * Kp_floor

   # Gravity:
   F += -9.81 * mass 
   y_dot = np.array([v, F], dtype=np.float)
   return y_dot

parameters = np.array([3, 50000, 3000], dtype=np.float)
y0 = np.array([1, 0], dtype=np.float) # Start at 1 meter, with 0 momentum

class SCIPY:
   def __init__(self, dt):
      self.dt = dt
   def name(self):
      return "scipy.integrate"
   def integrate(self, f, y0, parameters, endtime):
      timevector = np.arange(0.0, endtime, self.dt)
      states = scipy.integrate.odeint(f, y0, timevector, args=(parameters,))
      return timevector, states

class RK4:
   def __init__(self, dt):
      self.dt = dt

   def name(self):
      return "RK4"
   def integrate(self, f, y0, parameters, endtime):
      print(y0, parameters)
      t = 0.0
      states = []
      timevector = []
      state = y0
      states.append(state)
      timevector.append(t)
      while t < endtime:
         #### Runge Kutta integration:
         k1 = func(state,             t, parameters)
         k2 = func(state + 0.5*self.dt*k1, t, parameters)
         k3 = func(state + 0.5*self.dt*k1, t, parameters)
         k4 = func(state + self.dt*k3,     t, parameters)
         newstate = state + (self.dt/6.0)*(k1+2*k2+2*k3+k4)

         t += self.dt
         state = newstate

         states.append(state)
         timevector.append(t)

      return np.array(timevector, dtype=np.float), np.array(states, dtype=np.float)


class BDFfixed:
   def __init__(self, order):
      self.order = order
   def name(self):
      return "BDF"
   def integrate(self, f, y0, parameters, endtime):
      timevector = []
      states = []
      timevector.append(0.0)
      states.append(y0)
      t = 0.0
      h = 0.01
      state = y0
      while t < endtime:
         if len(states) < 4:
            k1 = func(state,             t, parameters)
            k2 = func(state + 0.5*h*k1, t, parameters)
            k3 = func(state + 0.5*h*k1, t, parameters)
            k4 = func(state + h*k3,     t, parameters)
            newstate = state + (h/6.0)*(k1+2*k2+2*k3+k4)
            state = newstate
            t += h
            timevector.append(t)
            states.append(state)
         else:
            t += h
            if np.max(state) > 10:
               break
            hyn = (11.0/6.0)*states[-1] + (-3)*states[-2] + (3/2)*states[-3] + (-1/3)*states[-4]
            state = state + hyn
            timevector.append(t)
            states.append(state)

      return np.array(timevector, dtype=np.float), np.array(states, dtype=np.float)

class WODE:
   def __init__(self):
      pass
   def integrate(self, f, y0, parameters, endtime):
      timevector = []
      states = []
      timevector.append(0.0)
      states.append(y0)
      return np.array(timevector, dtype=np.float), np.array(states, dtype=np.float)

class DOPRI45:
   def __init__(self):
      pass
   def name(self):
      return "DOPRI"
   def integrate(self, f, y0, parameters, endtime):
      timevector = []
      states = []
      timevector.append(0.0)
      states.append(y0)
      y = y0
      t = 0.0

      h = 0.01
      iteration = 0
      while t < endtime:
         k1 = h*f(y, t, parameters)
         k2 = h*f(y + (1.0/5.0)*k1, t + (1.0/5.0)*h, parameters)
         k3 = h*f(y + (3.0/40.0)*k1 + (9.0/40.0)*k2, t + (3.0/10.0)*h, parameters)
         k4 = h*f(y + (44.0/45.0)*k1 - (56.0/15.0)*k2 + (32.0/9.0)*k3, t + (4.0/5.0)*h, parameters)
         k5 = h*f(y + (19372.0/6561.0)*k1 - (25360.0/2187.0)*k2 + (64448.0/6561.0)*k3 - (212.0/729.0)*k4, t + (8.0/9.0)*h, parameters)
         k6 = h*f(y + (9017.0/3168.0)*k1 - (355.0/33.0)*k2 - (46732.0/5247.0)*k3 + (49.0/176.0)*k4 - (5103.0/18656.0)*k5, t + h, parameters)
         k7 = h*f(y + (35.0/384.0)*k1 + (500.0/1113.0)*k3 + (125.0/192.0)*k4 - (2187.0/6784.0)*k5 + (11.0/84.0)*k6, t + h, parameters)


         y_next = y + (35.0/384.0)*k1 + (500.0/1113.0)*k3 + (125.0/192.0)*k4 - (2187.0/6784.0)*k5 + (11.0/84.0)*k6
         z_next = y + (5179.0/57600.0)*k1 + (7571.0/16695.0)*k3 + (393.0/640.0)*k4 - (92097.0/339200.0)*k5 + (187.0/2100.0)*k6 + (1.0/40.0)*k7
         error = np.linalg.norm(y_next - z_next)
         #print(error, t)
         eps = 0.01
         if error < eps:
            y = y_next
            t = t + h
            timevector.append(t)
            states.append(y)

         s = ((eps*h)/(2*error)) ** (1.0/5.0)
         h = s * h
         if h < 1e-5:
            h = 1e-5
         iteration += 1
         if iteration > 10000:
            break

      return np.array(timevector, dtype=np.float), np.array(states, dtype=np.float)

rk4 = RK4(0.01)
sp = SCIPY(0.01)
wode = BDFfixed(4)
dopri = DOPRI45()
methods = [rk4, sp, dopri]

pylab.figure()
pylab.hold(True)
for method in methods:
   t, s = method.integrate(func, y0, parameters, 10)
   print(t.shape, s.shape)
   pylab.plot(t, s[:,0], 'x-', label=method.name())

pylab.legend()
pylab.show()