#include <iostream>
#include <vector>
#include <cmath>
#include <stdexcept>
#include "ODEfunc.h"

using namespace std;


typedef std::vector<double> Vec;
typedef Vec::value_type scalar;
class kepler2D : public ODEfunc<Vec> {
 public:
  // no private data, so trivial constructor
  kepler2D() { }

  // dim -- dimension of system
  int dim() { return 4; }

  // right-hand side for dq/dt = v, dv/dt = -q/||q||^3
  // format: x[0] = q[0], x[1] = q[1], x[2] = v[0], x[3] = v[1]
  Vec &eval(scalar t, const Vec &x, Vec &out)
  {
    if ( x.size() != 4 )
      throw runtime_error("Incorrect size: kepler2D");
    out.resize(x.size());
    out[0] = x[2];
    out[1] = x[3];
    scalar norm = sqrt(x[0]*x[0]+x[1]*x[1]);
    scalar denom = 1/(norm*norm*norm);
    out[2] = -x[0]*denom;
    out[3] = -x[1]*denom;
    return out;
  }
};

ostream &operator<<(ostream &s, Vec &z)
{
  s << "Vector:" << endl;
  for ( int i = 0; i < z.size(); ++i )
    s << z[i] << " ";
  s << endl;
  return s;
}

#define max(a,b)  ((a)<(b) ? (b) : (a))

scalar err_norm(Vec &x1, Vec &x2)
{
  scalar diff, max_diff;

  if ( x1.size() != x2.size() )
    throw runtime_error("Size mismatch: err_norm");
  max_diff = 0;
  for ( int i = 0; i < x1.size(); i++ )
    max_diff = max(max_diff,fabs(x1[i]-x2[i]));

  return max_diff;
}

int main(int argc, char *argv[])
{
  kepler2D myODE;
  Vec x0(myODE.dim()), xref(myODE.dim());
  scalar t0 = 0;
  scalar t;
  scalar t_end = 10.0;

  x0[0] = 1;
  x0[1] = 0;
  x0[2] = 0;
  x0[3] = 1;
  Vec x = x0;  // copy
  cout << "Initial vector:" << endl;
  cout << x;
  int n_steps;
  if ( argc > 1 )
    n_steps = atoi(argv[1]);
  else
    n_steps = 10;
  scalar h = (t_end-t0) / n_steps;
  cout << "Step size = " << h << endl;
  xref[0] = cos(t_end);
  xref[1] = sin(t_end);
  xref[2] = -sin(t_end);
  xref[3] = cos(t_end);

  euler(t0,x,myODE,h,n_steps);
  t = t0 + h*n_steps;
  cout << "At time t = " << t << " state is:" << endl;
  cout << x;
  cout << "Euler method error is " << err_norm(x,xref) << endl;

  x = x0;
  heun(t0,x,myODE,h,n_steps);
  t = t0 + h*n_steps;
  cout << "At time t = " << t << " state is:" << endl;
  cout << x;
  cout << "Heun method error is " << err_norm(x,xref) << endl;

  x = x0;
  rk4(t0,x,myODE,h,n_steps);
  t = t0 + h*n_steps;
  cout << "At time t = " << t << " state is:" << endl;
  cout << x;
  cout << "RK4 method error is " << err_norm(x,xref) << endl;

}