"""
ODE Solver: RK45 version 1.1.

Public domain, Connelly Barnes 2005.
"""

import math
import scipy
from scipy.linalg import norm
from scipy import array as vector

def solve(f, t0, tfinal, y0, tol = 1e-7):
  """
  Solve an ODE numerically using RK45.

  Solves dy/dt = f(t, y).  Returns a list of (t, y) tuples.
  Reference: http://www.library.cornell.edu/nr/bookcpdf/c16-2.pdf
  """

  def F(*args):
    return vector(f(*args))

  t = t0
  hmax = (tfinal - t0) / 128.0
  h = hmax / 4.0
  y = vector(y0)              # Column vector (nx1).
  out = [(t, list(y))]

  # Cash-Karp parameters
  a = [ 0.0, 0.2, 0.3, 0.6, 1.0, 0.875 ]
  b = [[],
       [0.2],
       [3.0/40.0, 9.0/40.0],
       [0.3, -0.9, 1.2],
       [-11.0/54.0, 2.5, -70.0/27.0, 35.0/27.0],
       [1631.0/55296.0, 175.0/512.0, 575.0/13824.0, 44275.0/110592.0, 253.0/4096.0]]
  c  = [37.0/378.0, 0.0, 250.0/621.0, 125.0/594.0, 0.0, 512.0/1771.0]
  dc = [c[0]-2825.0/27648.0, c[1]-0.0, c[2]-18575.0/48384.0,
        c[3]-13525.0/55296.0, c[4]-277.00/14336.0, c[5]-0.25]

  while t < tfinal:
    if t + h > tfinal:
      h = tfinal - t
    if t + h <= t:
      raise ValueError('Singularity in ODE')

    # Compute k[i] function values.
    k = [None] * 6
    k[0] = F(t, y)
    k[1] = F(t+a[1]*h, y+h*(k[0]*b[1][0]))
    k[2] = F(t+a[2]*h, y+h*(k[0]*b[2][0]+k[1]*b[2][1]))
    k[3] = F(t+a[3]*h, y+h*(k[0]*b[3][0]+k[1]*b[3][1]+k[2]*b[3][2]))
    k[4] = F(t+a[4]*h, y+h*(k[0]*b[4][0]+k[1]*b[4][1]+k[2]*b[4][2]+k[3]*b[4][3]))
    k[5] = F(t+a[5]*h, y+h*(k[0]*b[5][0]+k[1]*b[5][1]+k[2]*b[5][2]+k[3]*b[5][3]+k[4]*b[5][4]))

    # Estimate current error and current maximum error.
    E = norm(h*(k[0]*dc[0]+k[1]*dc[1]+k[2]*dc[2]+k[3]*dc[3]+k[4]*dc[4]+k[5]*dc[5]))
    Emax = tol*max(norm(y), 1.0)

    # Update solution if error is OK.
    if E < Emax:
      t += h
      y += h*(k[0]*c[0]+k[1]*c[1]+k[2]*c[2]+k[3]*c[3]+k[4]*c[4]+k[5]*c[5])
      out += [(t, list(y))]

    # Update step size
    if E > 0.0:
      h = min(hmax, 0.85*h*(Emax/E)**0.2)

  return out

def solve_func(f, t0, tfinal, y0, tol = 1e-5):
  """Returns y(t) function, uses linear interpolation."""
  L = solve(f, t0, tfinal, y0, tol)
  T = [t for (t, y) in L]
  Y = [y for (t, y) in L]
  interp_func = scipy.interpolate.interp1d(T, Y, axis=0)
  def func(t):
    return interp_func(t)[0]
  return func

def test_rk45(tol, max_error, max_evals):
  """
  Orbiting space shuttle test case.
  """

  y0 = vector([0.994, 0.0, 0.0, -2.00158510637908252240537862224])
  Period = 17.0652166
  fevals = [0]

  def yp(t, y):
    fevals[0] += 1
    mu = 0.012277471
    muhat = 1 - mu

    u1 = y[0]
    u1p = y[1]
    u2 = y[2]
    u2p = y[3]

    D1 = ((u1 + mu)**2 + u2**2)
    D1 = D1 * math.sqrt(D1)
    D2 = ((u1 - muhat)**2 + u2**2)
    D2 = D2 * math.sqrt(D2)

    return (u1p,
            u1 + 2 * u2p - muhat*(u1+mu) / D1 - mu*(u1-muhat) / D2,
            u2p,
            u2 - 2 * u1p - muhat * u2 / D1 - mu * u2 / D2)
  
  sol = solve(yp, 0.0, Period, y0, tol)
  ylast = sol[len(sol)-1][1]

#  print 'Fevals:', fevals[0]
#  print 'Error:', norm(ylast - y0)
  assert fevals[0] < max_evals
  assert norm(ylast - y0) < max_error

def test():
  """
  Unit tests.
  """

  print 'Testing:'
  test_rk45(tol=1e-6, max_error=0.01, max_evals=2000)
  test_rk45(tol=1e-9, max_error=3e-5, max_evals=5000)
  print '  rk45:        OK'

if __name__ == '__main__':
  test()