// Planet2.cc
//  see also gnu20, gnu21
// usage:
//  Planet2 1300 0.03 25 | gnuplot -persist
// This program requires the directory 'output' to exist, and
// expects it to contain a file gnu0.

#include <iostream>
#include <fstream>
#include <cmath>
#include <cstdlib>
using namespace std;

#define D  2  // number of dimensions
struct particle {
  double x[D] ; // (x,y) coordinates
  double p[D] ; // momentum
  double F[D] ; // force on particle
  double im   ; // inverse mass of particle
  double GMm  ; // gravitational parameter of particle
  double v[D] ; // velocity
  double T    ; // kinetic energy
  double J    ; // angular momentum
  double V    ; // potential energy
  double r    ; // absolute distance from origin
  int  law    ; // which force law
} ; // Note the definition of a structure ends with a semicolon

struct control {
  int verbose    ; // program verbosity
  int printing   ; // period with which to print
  int commanding ; // period with which to issue commands
  //  ofstream *fout ; // where we write data to
  int time_in_ms ; // time between frames (commands)
  clock_t latest_time;
  ofstream fout  ; 
  int euler      ; // whether to do euler method after leapfrog
};

// the "&" in "&a" indicates 'this is a pass by reference'. 
void Force( particle &a )
  // sets the force vector and the potential energy
{
  double R = 0.0 ; 
  for ( int i = 0 ; i < D ; i++ ) {
    R += a.x[i] * a.x[i] ;
  }
  double r = sqrt(R) ;
  a.r = r ; 
  if ( a.law == -2 ) { // inverse square law
    a.V = - a.GMm / r ;
    for ( int i = 0 ; i < D ; i++ ) {
      a.F[i] = (- a.GMm * a.x[i] / (r*R) ) ;  // inverse sq
    }
  } else { // Hooke Law 
    a.V = 0.5 * a.GMm * R ; 
    for ( int i = 0 ; i < D ; i++ ) {
      a.F[i] = (- a.GMm * a.x[i])  ;  // linear spring law
    }
  }
}

// Implement     x += v dt
void PositionStep ( particle &a , double dt )
{ 
  for ( int i = 0 ; i < D ; i++ ) 
    a.x[i] += dt * a.p[i] * a.im ; 
}

// Implement     p += f dt
void MomentumStep ( particle &a , double dt )
{ 
  for ( int i = 0 ; i < D ; i++ ) 
    a.p[i] += dt * a.F[i] ;
}

void v2p( particle &a )
  // propagate the changed velocity into the momentum vector
{
  for ( int i = 0 ; i < D ; i++ ) 
    a.p[i] =  a.v[i] / a.im ;
}

void p2v ( particle &a ) {
  for ( int i = 0 ; i < D ; i++ ) 
    a.v[i]  = a.p[i] * a.im ;
}

void pv2T ( particle &a ) {
  // compute the kinetic energy and angular momentum
  a.T=0.0;
  for ( int i = 0 ; i < D ; i++ ) { 
    a.T    += 0.5*a.v[i] * a.p[i] ; 
  }
  a.J  =  a.x[0] * a.p[1] - a.x[1] * a.p[0] ; 
}

// Print out the state (x,v,Energies,J) in columns
void showState ( particle &a , ostream &fout )
{
  for ( int i = 0 ; i < D ; i++ ) { 
    fout << "\t"<<a.x[i];
  }
  for ( int i = 0 ; i < D ; i++ ) { 
    fout << "\t"<<a.v[i];
  }
  fout << "\t" << a.T     // KE
       << "\t" << a.V     // PE
       << "\t" << a.T+a.V // Total E
       << "\t" << a.J     // Angular momentum
       << endl;
  fout << "#" ;
  for ( int i = 0 ; i < D ; i++ ) { 
    fout << "\tx["<<i<<"]\t";
  }
  for ( int i = 0 ; i < D ; i++ ) { 
    fout << "\tv["<<i<<"]\t";
  }
  fout << "\tT\t\tV\t\tT+V\t\tJ" ; 
  fout << endl;
}

void showStateByArrow ( particle &a , ostream &fout )
{
  fout << "set arrow 101 from " << a.x[0] << ","
       << a.x[1]  << " to "
       << a.x[0]+a.v[0] << ","
       << a.x[1]+a.v[1] << " lt 5 " << endl ; 
}

void PossibleOutput( particle &a , control &c , int i , double t ) {
  if( ( c.printing && !(i%c.printing) )
      || ( c.commanding  && !(i%c.commanding) ) ) {
    p2v(a) ;  pv2T(a) ;  Force( a ) ;
  }
  if( ( c.printing && !(i%c.printing) ) ) {
    c.fout << t ; showState( a ,  c.fout ) ;
  }
  if ( c.commanding  && !(i%c.commanding) ) {
    // SEND COMMANDS TO gnuplot live via 'cout' pipe
    c.fout.flush() ; // make sure file is out
    system("tail -39 /tmp/1 > /tmp/2"); 
    showStateByArrow( a , cout ) ; // sends 'set arrow' command to gnuplot
    if(i==0) {
      cout << "load 'output/gnu0'" << endl ; // to gnuplot
      cout.flush() ;
      c.latest_time = clock() ;
    } else {
      c.latest_time += CLOCKS_PER_SEC * c.time_in_ms / 1000 ;
      clock_t endwait = c.latest_time ;
      clock_t time_to_go = endwait-clock() ; 
      while( clock() < endwait ) {} ;       // wait till the appointed time
      cout << "load 'output/gnu0'" << endl ; // to gnuplot
      //      cout << "replot" << endl ;
      while( clock() < endwait+time_to_go/3 ) {} ; // add a little further
      // pause to give gnuplot a chance to do its thing before we edit
      // files again
    }
  }
}

// The Leapfrog method simulates the equations of motion
//  dx/dt = p/m
//  dp/dt = f
// by alternately updating the position (using the latest
// momentum), then the momentum (using the force at the
// updated position).
void LeapfrogDynamics( particle &a, 
		       double dt  , // step size
		       double &t  , // current time
		       int N      , // number of steps
		       control &c   // parameters controlling output, etc
		       )
{
  for ( int i=0 ; i < N ; i ++ ) {
    PossibleOutput( a , c , i , t ) ; 
    // each iteration we move the position a half-step
    PositionStep( a , dt*0.5 ) ;
    Force( a ) ; // (computes the force at this position)
    // then we move the momentum full-step
    MomentumStep( a , dt )     ;
    // then we move the position a half-step
    PositionStep( a , dt*0.5 ) ;
    t += dt ;
  }
}

// The Euler method updates the position and momentum
// simultaneously.  For the force laws considered here,
// we can get the effect of simultaneous update by
// finding the force, then 
// updating the position, THEN the momentum.
void EulerDynamics( particle &a, 
		    double dt  , // step size
		    double &t  , // current time
		    int N      , // number of steps
		    control &c   // parameters controlling output, etc
		    )
{
  for ( int i=0 ; i < N ; i ++ ) {
    PossibleOutput( a , c , i , t ) ; 
    Force( a ) ; // (computes the force at this position)
    PositionStep( a , dt ) ;
    MomentumStep( a , dt ) ;
    t += dt ;
  }
}

int main(int argc, char* argv[])
{
  particle   a ;
  control    c ;
  // set defaults
  int N = 250 ;  
  double dt = 0.1 ;
  double t = 0.0 ;
  c.verbose = 1 ;
  c.printing   =  4 ;
  c.commanding =  1 ; // number of iterations between plots
  c.time_in_ms = 100; // real time between plots (in ms)
  c.euler = 0 ;       // whether to do euler after leapfrog

  // read in any command-line arguments
  if(argc>1)   {
    sscanf( argv[1], "%d", &N ) ; // put the first command-line argument in N
  }
  if(argc>2) {
    sscanf( argv[2], "%lf", &dt ) ; // put the 2nd argument in dt
  }
  if(argc>3) {
    sscanf( argv[3], "%d", &(c.time_in_ms) ) ;
  }

  // try to write output file
  char fileName[]="/tmp/1";
  c.fout.open(fileName);
  if(c.fout.good()==false){
    cerr << "can't write to file " << fileName << endl;
    exit(0);
  }

  // set initial conditions for a particle
  a.law  = -2 ; // -2: inverse square
  a.im   = 1.0 ;
  a.v[0] = 0.4;
  a.v[1] = 0.0;
  a.x[0] = 1.5;
  a.x[1] = 2;
  a.GMm = 1.0;
  v2p(a) ; // get the initial momentum from the velocity

  LeapfrogDynamics( a , dt , t , N , c ) ;
  if ( c.euler ) {
    cerr << "# SWITCHING TO EULER\n" ;
    c.fout << "\n\n\n" ; 
    EulerDynamics( a , dt , t , N , c ) ;
  }
  
  return 0;
}