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

using namespace std;

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

double RAD_PLUMMER = 0.0;

// Setup the solar system: Now we need to setup several planets which we collect into a 'vector'
void initializeSolSystem(vector <state> &start) {
#ifdef SELF_GRAV
    start.push_back(state(0/l_unit        , 0/v_unit      , msun/m_unit));       // sun
#endif
    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 initialize2(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 initializePerl(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 calculateEffectivePotential(state &here) {
    vec L = here.x.cross(here.v);
    double pot = G * here.m * (msun / m_unit) / 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 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(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 * (msun / m_unit) / r / r / r);
    }
}

// For task 2: self gravity
void calculateAccelerationSelf(vector <state> &here) {
    for (int i = 0; i < here.size(); i++) {
        here[i].a = vec();
    }
    for (int i = 0; i < here.size(); i++) {
        for (int j = i + 1; j < here.size(); j++) {
            vec dist_vec = here[i].x - here[j].x;
            double r = sqrt(dist_vec.abs2() + RAD_PLUMMER * RAD_PLUMMER);
            vec apm = -(G / (r * r * r)) * dist_vec;

            here[i].a += apm * here[j].m;
            here[j].a -= apm * here[i].m;
        }
    }
}

// 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) {
    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:
            initializeSolSystem(current);
            RAD_PLUMMER = 0.0;         // if we set RAD_PLUMMER to 0 then it behaves like a point mass
#ifdef SELF_GRAV
            RAD_PLUMMER = 0.01 * au / l_unit;
#endif
            outfile.open("sol_system.dat");
            break;
        case 2:
            initialize2(current);
            RAD_PLUMMER = (0.5 * au / l_unit);
            outfile.open("plummer_2.dat");
            break;
        case 3:
            initializePerl(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;

    // Calculate initial state acceleration

#ifdef SELF_GRAV
    calculateAccelerationSelf(current);
#else
    calculateAcceleration(current);
#endif

    // This is loop which is done until integration of t reaches tmax
    while (current[0].t < TMAX) {
        // Write output (only if time defined by OUTPUT_FREQUENCY is reached !)
        if (current[0].t > next_output_time) {
            next_output_time += OUTPUT_FREQENCY;
            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 << " "
                        << calculateEffectivePotential(current[i]) << " "
                        << calculateTotalEnergy(current[i]) << " ";
            outfile << endl;
        }

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


        // KDK Scheme
        kick(current, 0.5 * dt);
        drift(current, 1.0 * dt);
#ifdef SELF_GRAV
        calculateAccelerationSelf(current);
#else
        calculateAcceleration(current);
#endif
        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;
}