#!/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

# ODE integrators
def EulerStep(u,t,dt,rhs):
  up=u + dt*rhs(u, t)
  return up

def RK2Step(u,t,dt,rhs):
  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):
  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


# spatial finite differences
def FD1(u,h): # 1st derivative
  up = np.zeros(len(u))
  up[1:-1] =  (u[2:] - u[:-2]) / (2.0*h) # interior of grid
  # compute derivative at boundaries:
  up[ 0] = (    -u[ 2] + 4.0*u[ 1] - 3.0*u[0]) / (2.0*h);
  up[-1] = (-3.0*u[-1] + 4.0*u[-2] - u[-3])    / (2.0*h) # 2nd order accurate
  return up

def FD1periodic(u,h): # 1st derivative with periodic BCs
  up = np.zeros(len(u))
  up[1:-1] =  (u[2:] - u[:-2]) / (2.0*h) # interior of grid
  # compute derivative at boundaries: identify points 0 and -1
  up[ 0] = ( u[1] - u[-2] ) / (2.0*h);
  up[-1] = ( u[1] - u[-2] ) / (2.0*h) # 2nd order accurate
  return up


def FD2periodic(u,h): # 2nd derivative with periodic BCs
  ddu = np.zeros(len(u))
  ddu[1:-1] = (u[2:] - 2.0*u[1:-1] + u[:-2]) / h**2
  ddu[0]    = (u[1] - 2.0*u[ 0] + u[-2]) / h**2
  ddu[-1]   = (u[1] - 2.0*u[-1] + u[-2]) / h**2
  return ddu

def FD2(u,h): # 2nd derivative at interior gridpoints
  ddu = np.zeros(len(u))
  ddu[1:-1] = (u[2:] - 2.0*u[1:-1] + u[:-2]) / h**2
  # boundary terms missing!
  return ddu
  
  

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

  gridlen = len(u0[0])

  [tmax, nstepsmax, nequations, npoints, deltax, tsample, xsample] = param
  nsteps = nstepsmax # we don't really know how many steps we'll need!
  
  t = 0.0  
  n = 0 
  
  u = u0
  while n <= nsteps-1 and t < tmax: 
    up = ODEstepper(u,t,dt,ODErhs)
    u = up

    if np.mod(n,tsample) == 0:
      for component in range(nequations):
        ODEsave(solnames[component], t, u[component], param)

    n+=1
    t+=dt
    
  print 'stopped at n = %d, t = %.4f' %(n, t)
    
  return (t, u)

def ODEsave(name, t, sol, param):

  [tmax, nstepsmax, nequations, npoints, deltax, tsample, xsample] = param
  
  fname = name + '.dat' 
  f_handle = file(fname, 'a')

#  coordvect = grid = np.arange(npoints) * deltax
#  ts = np.column_stack([coordvect, sol])
  ts = sol[::xsample]

  np.savetxt(f_handle, ts)
  f_handle.write('\n\n')
  f_handle.close()