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

using namespace std;

#include "const.h"
#include "1sample.h"


// 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);
#ifndef EXCENTRIC
    start.v = vec(0, 2.0 * M_PI * au / (sec_per_year * v_unit), 0);
#else
    start.v = vec(0, 0.7 * 2.0 * M_PI * au / (sec_per_year * v_unit), 0);
#endif
}

/* 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 calculateTotalEnergy(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 */
}

void calculateAcceleration(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 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;

#ifdef DKD
    cout << "# Integrating with DKD-scheme" << endl;
#else
    cout << "# Integrating with KDK-scheme" << endl;
#endif
#ifdef VAR_TIMESTEP
    cout << "# Using variable timesteps" << endl;
#endif

    // Time integration parameter
    int orbit_counter = 0;
    double dt = DT_FACTOR;
    vec x0;

    // Lets initialize the system. We defined this procedure before
    initialize(current);

    // Calculate initial state acceleration
    calculateAcceleration(current);

    // This is loop which is done until integration of t reaches tmax
    const double torbit = sec_per_year / t_unit;
    const double tmax = NOrbits * torbit;
    while (current.t < tmax) {
        // Lets output time , position and total energy of the current time
        if (orbit_counter % OUTPUT_FREQUENCY == 0) {
            cout << current.t << current.x << current.v << calculateTotalEnergy(current) << endl;
        }

        x0 = current.x;

        // Variable Timestep
#ifdef VAR_TIMESTEP
        dt = DT_FACTOR / current.a.abs();
#endif

#ifdef DKD
        // DKD Scheme
        drift(current, 0.5 * dt);
        calculateAcceleration(current);
        kick(current, 1.0 * dt);
        drift(current, 0.5 * dt);
#else
        // KDK Scheme
        kick(current, 0.5 * dt);
        drift(current, 1.0 * dt);
        calculateAcceleration(current);
        kick(current, 0.5 * dt);
#endif

        // 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;
}