#include <iostream>
#include <cmath>
#include <vector>

// We can set TMAX either here or in the Makefile 
#ifndef TMAX
#define TMAX 20
#endif

#ifndef OUTPUT_FREQENCY
#define OUTPUT_FREQENCY 40
#endif

using namespace std;

#include "example.h"
#include "const.h"

// Time integration parameter
const double dt_factor=0.005;

// For time variable steps
const double dt_factor_time_dependent=0.5;

// Setup the earth orbit (1 au, 1 year period) for the state which we start with
void initialize(state &start)
{
  start.t = 0;
  start.m = mearth / m_unit;  
  start.x = vec(au / l_unit,0,0);
  start.v = vec(0,2.0 * M_PI,0);
  // Eccentric orbit
  // start.v = vec(0,0.8 * 2.0 * M_PI,0);
}

/* Calculate total energy at a given state. This is a function which gets a state as an argument 
   and returns a double (e.g. the total energy of the state) */
double calculate_total_energy(state &here)
{
  double mass = msun / m_unit;                   /* Mass of sun in internal units 
                                                    (will be 1 in the chosen units) */
  double pot = here.m * mass * G / here.x.abs(); /* Potential energy */
  double kin = 0.5 * here.m * here.v.abs2();     /* Kinetic energy */

  return(kin-pot);                               /* Here we return the total energy */
}

// Calculate acceleration of a given state. This is a procedure which gets a state as an argument
void calculate_acceleration(state &here)
{
  double r = here.x.abs();
  double sol_mass = msun / m_unit;               /* Mass of sun in internal units 
						    (will be 1 in the chosen units) */
  here.a = here.x * (-G * sol_mass / r / r / r);
}

// update velocities of a given state. This is a procedure which gets a state and a deltaT as arguments
void kick(state &here, double mydt)
{
  here.v += here.a * mydt;
}

// update positions of a given state. This is a procedure which gets a state and a deltaT as argument
void drift(state &here, double mydt)
{
  here.x += here.v * mydt;
}


int main (int argc, const char **argv)
{
  /* state is a structure (we defined it in example.h) which contains all quantities (like position, velocity etc)
     needed to describe a test particle.
     We want to create one instance of it which we call current for or actual time step */
  
  state current;

  // Lets output at the beginning some useful information ... 
  cout << "# L = " << l_unit << endl;
  cout << "# T = " << t_unit << endl;
  cout << "# M = " << m_unit << endl;
  cout << "# V = " << v_unit << endl;
  cout << "# G = " << G << endl;

// Time integration parameter
  int orbit_counter=0;
  double dt = 1.0 * dt_factor;
  vec x0;
  
  // Lets initialize the system. We defined this procedure before 
  initialize(current);

  // Now we calculate the acceleration of the initial state
  calculate_acceleration(current);

  // This is loop which is done until integration of t reaches our defined TMAX
  while(current.t < TMAX)
    {
      // Lets output time , position and total energy of the current time
      if( orbit_counter % OUTPUT_FREQENCY == 0)
	{
	  cout << current.t << current.x << current.v << calculate_total_energy(current) << endl;
	}

      x0 = current.x;

      // Variable Timestep
      // dt = dt_factor_time_dependent / current.a.abs();

      // KDK Scheme
      kick(current, 0.5 * dt);
      drift(current, 1.0 * dt);
      calculate_acceleration(current);
      kick(current, 0.5 * dt);

      // DKD Scheme 
      /*
      drift(current, 0.5 * dt);
      calculate_acceleration(current);
      kick(current, 1.0 * dt);
      drift(current, 0.5 * dt);
      */

      // This is a trick to count the orbits (as sign reversal from - to + in the x position 
      if(x0.comp[0] < 0 && current.x.comp[0] >=0)
	orbit_counter++;

      // Finally we update the time t
      current.t += dt;
    }

  return(0);
}


/* How to get it running:
  1) make
  2) ./myprogram >& orbit.dat
  3) gnuplot
     gnuplot> plot 'orbit.dat' u ($2):($3)
*/