#include <iostream>
#include <fstream>
#include <cmath>
#include <stdlib.h>
#include <vector>

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

#ifndef OUTPUT_FREQENCY
#define OUTPUT_FREQENCY 0.01
#endif

#define NTEST 100
#define MASS (msun / m_unit)

using namespace std;

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

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

double RAD_PLUMMER=0;

// Setup the solar system: Now we need to setup several planets which we collect into a 'vector'
void initialize_sol_system(vector<state> &start)
{
  start.push_back(state(0.3871*au/l_unit, 47.87e5/v_unit, 3.302e26/m_unit));   // merkur
  start.push_back(state(0.723 *au/l_unit, 35.02e5/v_unit, 4.869e27/m_unit));   // venus
  start.push_back(state(1.0   *au/l_unit, 29.78e5/v_unit, 5.97219e27/m_unit)); // earth
  start.push_back(state(1.524 *au/l_unit, 24.13e5/v_unit, 6.219e26/m_unit));   // mars
  start.push_back(state(5.203 *au/l_unit, 13.07e5/v_unit, 1.899e30/m_unit));   // jupiter
  start.push_back(state(9.5826*au/l_unit,  9.69e5/v_unit, 5.685e29/m_unit));   // saturn
  start.push_back(state(19.201*au/l_unit,  6.81e5/v_unit, 8.683e29/m_unit));   // uranus
  start.push_back(state(30.047*au/l_unit,  5.43e5/v_unit, 1.0243e29/m_unit));  // neptun
  start.push_back(state(39.482*au/l_unit,  4.72e5/v_unit, 1.25e25/m_unit));    // pluto
}

// Setup two particles in potential
void initialize_2(vector<state> &start)
{
  start.push_back(state(au/l_unit, 0.02*29.78e5/v_unit, 1));
  start.push_back(state(au/l_unit, 0.1*29.78e5/v_unit, 1));
}

// Setup perl of test particles
void initialize_perl(vector<state> &start)
{
  for(int i=0;i<NTEST;i++)
    {
      state test(au / l_unit, 29.78e5/v_unit*(float)i/ NTEST, 1);
      test.x.comp[1] = 0.0 + (float)i / NTEST / 10 * au / l_unit;
      start.push_back(test);
    }
}

// Calculate effective potential
double calculate_effective_potential(state &here)
{
  vec L=here.x.cross(here.v);
  double pot = G * here.m * MASS / sqrt(here.x.abs2() + RAD_PLUMMER * RAD_PLUMMER); /* Potential energy */
  double red = 0.5 * here.m * L.abs2() / here.x.abs2();
  return(red-pot);
}

/* 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 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(vector<state> &here)
{
  for(int i=0;i<here.size();i++)
    {
      double r = sqrt(here[i].x.abs2() + RAD_PLUMMER * RAD_PLUMMER);
      here[i].a = here[i].x * (-G * 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(vector<state> &here, double mydt)
{
  for(int i=0;i<here.size();i++)
    here[i].v += here[i].a * mydt;
}

// update positions of a given state. This is a procedure which gets a state and a deltaT as argument
void drift(vector<state> &here, double mydt)
{
  for(int i=0;i<here.size();i++)
    here[i].x += here[i].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 */
  
  vector<state> current;          // Vector of our particle system
  double next_output_time = 0;    // Time for next output
  double dt=0;                    // Timestep to be defined later
  ofstream outfile;               // Lets use directly a file to output the data

  // Lets check if an argument was passed to out program
  if(argc<2)
    {
      cout << "please give a number to the program:" << endl << "1: Sol System" << endl
	                                                     << "2: 2 Particles" << endl
	                                                     << "3: Particle chain" << endl;
      return(0);
    }

  // Lets initialize the system depending on users choice 
  switch(atoi(argv[1]))
    {
    case 1:
      initialize_sol_system(current);  
      RAD_PLUMMER =0;                  // if we set RAD_PLUMMER to 0 then it behaves like a point mass
      outfile.open("sol_system.dat");
      break;
    case 2:
      initialize_2(current);
      RAD_PLUMMER = (0.5 * au / l_unit);
      outfile.open("plummer_2.dat");
      break;
    case 3:
      initialize_perl(current);
      RAD_PLUMMER = (0.5 * au / l_unit);
      outfile.open("plummer_perl.dat");
      break;
    default:
      cout << "Only give 1,2 or 3 !" << endl;
      return(0);
    }
  
  // 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;

  // 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[0].t < TMAX)
    {
      // Write output (trick: do it only if time defined by OUTPUT_FREQUENCY is reached !)
      if(current[0].t > next_output_time)
	{
	  cout << current[0].t << " " << dt << " " << next_output_time << endl;

	  outfile << current[0].t << " ";
	  for(int i=0; i< current.size(); i++)
	    outfile << current[i].x << current[i].v << " " 
		 << calculate_effective_potential(current[i]) << " "
		 << calculate_total_energy(current[i]) << " ";
	  outfile << endl;
	  next_output_time += OUTPUT_FREQENCY;
	}

      // Variable Timestep (trick: first find largest acceleration and convert to timestep at the end !)
      double a_max=0;
      for(int i=0; i< current.size(); i++)
	if(current[i].a.abs() > a_max)
	  a_max = current[i].a.abs();
      dt = dt_factor_time_dependent / a_max;


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

      // Finally we update the time t
      for(int i=0; i< current.size(); i++)
	current[i].t += dt;

    }

  outfile.close();
  
  return(0);
}


/* How to get it running:
  1) make
  2) ./myprogram 
*/