#!/usr/bin/python

# (C) Michael Puerrer, Sascha Husa

import matplotlib.pyplot as plt
import numpy as np

from math import pi

from scipy.interpolate import UnivariateSpline
from scipy import optimize

#################

def ODErhsExample(u,t):
  [x, y] = u

  XRHS = -x
  YRHS = -y
    
  return np.array([XRHS, YRHS])
  
#################


# ODE integrators
def EulerStep(u,t,dt,rhs):
  n=len(u)
  up=np.zeros(n)
  up=u + dt*rhs(u, t)
  return up

def RK2Step(u,t,dt,rhs):
  n=len(u)
  up=np.zeros(n)
  k1=np.zeros(n)
  k2=np.zeros(n)
  
  k1 = dt*rhs(u,t)
  k2 = dt*rhs(u + k1, t + dt)
  
  up = u + 0.5*(k1 + k2)
  return up

def RK4Step(u,t,dt,rhs):
  n=len(u)
  up=np.zeros(n)
  k1=np.zeros(n)
  k2=np.zeros(n)
  k3=np.zeros(n)
  k4=np.zeros(n)

  k1 = dt*rhs(u,t)
  k2 = dt*rhs(u + 0.5*k1, t + 0.5*dt)
  k3 = dt*rhs(u + 0.5*k2, t + 0.5*dt)
  k4 = dt*rhs(u + k3, t + dt)

  up = u + (k1 + 2*k2 + 2*k3 + k4) / 6.0
  return up



# ODE integration driver
def ODEint(u0, dt, param, ODEstepper, ODErhs):

  [tmax, nstepsmax, nequations] = param
  nsteps = nstepsmax # we don't really know how many steps we'll need!
  sol = np.zeros((nsteps,nequations)) # to store the solution
  
  t=0.0  
  u = u0

  n=0  
  sol[0] = u
  while n <= nsteps-1 and t < tmax: 
    u = ODEstepper(u,t,dt,ODErhs)
    n+=1
    sol[n] = u
    t+=dt
    
  return (t, sol[:n])

def ODEsave(name, tEnd, sol, solnames):
  coordvect = np.linspace(0.0,tEnd,len(sol))
  
  ts = np.column_stack([coordvect, (sol.transpose())[0]])

  np.savetxt(name + solnames[0] + '.dat', ts, fmt='%.18e', delimiter=' ', newline='\n')


def ODEplot(tEnd, sol, solnames):

  components = len(sol[1])
  timesteps  = len(sol)
  
  coordvect = np.linspace(0.0, tEnd, timesteps)

  plt.figure(1)
  for component in range(components):
	plt.subplot(2,1,component)
	values = sol.transpose()
  	plt.plot(coordvect,values[component],'b-')
  	plt.xlabel('t')
  	plt.ylabel(solnames[component])
  	plt.title(solnames[component])
  
  plt.show()
  
  
if __name__ == "__main__":  # now we can import this file as a module without running 'main', or execute it as a script
  
  nequations = 3
  u0 = [5., 6.]
  tmax = 1
  nstepsmax = 100000
  param =[tmax, nstepsmax, nequations]

  
  dt0 = 0.005
  
  dt1 = dt0 # set stepsize
  (tE1, sol1) = ODEint(u0, dt1, param)
  print 'stopped at n = %d, t = %.4f, state vector = ' %(len(sol1), tE1) #, sol1[-1])   
  ODEsave('dt1', tE1, sol1, ['x','y'])
  ODEplot(tE1, sol1, ['x','y'])

  dt2 = dt0/2 # set stepsize
  (tE2, sol2) = ODEint(u0, dt2, param)
  print 'stopped at n = %d, t = %.4f, state vector = ' %(len(sol2), tE2) #, sol2[-1])   
  ODEsave('dt2', tE2, sol2, ['x','y'])
  ODEplot(tE2, sol2, ['x','y'])

  dt3 = dt0/4 # set stepsize
  (tE3, sol3) = ODEint(u0, dt3, param)
  print 'stopped at n = %d, t = %.4f, state vector = ' %(len(sol3), tE3) #, sol3[-1])   
  ODEsave('dt3', tE3, sol3, ['x','y'])
  ODEplot(tE3, sol3, ['x','y'])


  diff12 = sol1                    - (sol2[::2])[:len(sol1)]
  diff23 = (sol2[::2])[:len(sol1)] - (sol3[::4])[:len(sol1)]

  ODEsave('diff12', tE1, diff12, '1-2')
  ODEsave('diff23', tE1, diff23, '2-3')