// Vec is a vector class with scalar type T
// This is an abstract (template) class
template<class Vec, class T=typename Vec::value_type>
class ODEfunc {
 public:
  // f(t,x) is put in out, which is returned
  virtual Vec &eval(T t, const Vec &x, Vec &out) = 0;
  virtual int dim() = 0;
};

// Euler method -- carries out n_step steps of Euler's method
// -- x is the initial state on entry, final state on exit returned)
// -- t0 is the initial time
// -- h is the step size
template<class Vec, class T>
Vec &euler(T t0, Vec &x, ODEfunc<Vec,T> &f, T h, int n_steps)
{
  Vec temp(x.size());
  T   t;

  t = t0;

  for ( int i = 0; i < n_steps; ++i )
    {
      temp = f.eval(t,x,temp);
      for ( int j = 0; j < x.size(); ++j )
        x[j] = x[j] + h*temp[j];
      t = t + h;
    }

  return x;
}

// Heun's method -- 2nd order
// -- carries out n_step steps of this method
// -- x is the initial state on entry, final state on exit (returned)
// -- t0 is the initial time
// -- h is the step size
template<class Vec, class T>
Vec &heun(T t0, Vec &x, ODEfunc<Vec,T> &f, T h, int n_steps)
{
  Vec temp(x.size());
  Vec v1(x.size()), v2(x.size());
  T   t;

  t = t0;
  for ( int i = 0; i < n_steps; ++i )
    {
      f.eval(t,x,v1);
      for ( int j = 0; j < x.size(); ++j )
        temp[j] = x[j] + h*v1[j];
      f.eval(t+h,temp,v2);
      for ( int j = 0; j < x.size(); ++j )
        x[j] = x[j] + (h/2.0)*(v1[j]+v2[j]);
      t = t + h;
    }

  return x;
}


// Standard 4th order Runge-Kutta method
// -- carries out n_step steps of this method
// -- x is the initial state on entry, final state on exit (returned)
// -- t0 is the initial time
// -- h is the step size
template<class Vec, class T>
Vec &rk4(T t0, Vec &x, ODEfunc<Vec,T> &f, T h, int n_steps)
{
  Vec temp(x.size());
  Vec v1(x.size()), v2(x.size()), v3(x.size()), v4(x.size());
  T   t;

  t = t0;
  for ( int i = 0; i < n_steps; ++i )
    {
      f.eval(t,x,v1);
      for ( int j = 0; j < x.size(); ++j )
        temp[j] = x[j] + 0.5*h*v1[j];
      f.eval(t+0.5*h,temp,v2);
      for ( int j = 0; j < x.size(); ++j )
        temp[j] = x[j] + 0.5*h*v2[j];
      f.eval(t+0.5*h,temp,v3);
      for ( int j = 0; j < x.size(); ++j )
        temp[j] = x[j] + h*v3[j];
      f.eval(t+h,temp,v4);
      for ( int j = 0; j < x.size(); ++j )
        x[j] = x[j] + (h/6.0)*(v1[j]+2*v2[j]+2*v3[j]+v4[j]);
      t = t + h;
    }

  return x;
}