#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#include "nbgenerator.hh"
#include "nbio.hh"
#include "nbtypes.hh"
#include "nbstopwatch.hh"

/*
 * Global constants
 */

/* set gravity constant to one */
const double G = 1.0;
/* The epsilon^2 for the softening */
const double epsilon2 = 1E-14;

/**
 * Compute acceleration
 *
 * \param[in]  n Number of bodies
 * \param[in]  r Positions of the bodies
 * \param[in]  m Masses of the bodies
 * \param[out] a The computed accelerations of the bodies
 *
 * The inner loop executes 26 floating point operations and is executed
 * ((n*n)-n)/2 times, resulting in 13*n*(n-1) floating point operations total
 * per call to acceleration().
 */
void acceleration (int n, vector3 r[], double m[], vector3 a[])
{
  int i,j;
  double d0,d1,d2,d,d_sq,factori,factorj;
  double invfact;

  /* compute acceleration exploiting symmetry */
  for (i=0; i<n; i++)
    for (j=i+1; j<n; j++) {
      /* compute distance vector */
      d0 = r[j][0]-r[i][0];
      d1 = r[j][1]-r[i][1];
      d2 = r[j][2]-r[i][2];

      /* compute (scalar valued) distance and distance^2, with softening */
      d_sq = d0*d0 + d1*d1 + d2*d2 + epsilon2;
      d = sqrt(d_sq);

      /* precompute the factors in the acceleration equation */
      invfact = G/(d*d_sq);
      factori = m[i]*invfact;
      factorj = m[j]*invfact;

      /* compute acceleration of body i due to body j */
      a[i][0] += factorj*d0;
      a[i][1] += factorj*d1;
      a[i][2] += factorj*d2;

      /* compute acceleration of body j due to body i */
      a[j][0] -= factori*d0;
      a[j][1] -= factori*d1;
      a[j][2] -= factori*d2;
    }
}

/**
 * Do one time step
 *
 * \param[in]     n    Number of bodies
 * \param[in]     dt   Time step size
 * \param[in,out] x    IN: old positions, OUT: new positions
 * \param[in,out] v    IN: old velocities, OUT: new velocities
 * \param[in]     m    Masses
 * \param[in]     a    accelerations
 *
 * The "update position" and "update velocity" execute 6 floating point
 * operations per iteration each, resulting in 12*n floating point ops
 * per call to euler().
 */
void euler (const int N, double dt, vector3 r[], vector3 v[], double m[], vector3 a[])
{
    int i;
    
    /* update position and velocity */
    for (i=0; i<N; i++)
    {
        r[i][0] += dt*v[i][0];
        r[i][1] += dt*v[i][1];
        r[i][2] += dt*v[i][2];

        v[i][0] += dt*a[i][0];
        v[i][1] += dt*a[i][1];
        v[i][2] += dt*a[i][2];
    }
}


/**
 * Do one time step
 *
 * \param[in]     n    Number of bodies
 * \param[in]     dt   Time step size
 * \param[in,out] x    IN: old positions, OUT: new positions
 * \param[in,out] v    IN: old velocities, OUT: new velocities
 * \param[in]     m    Masses
 * \param[in,out] a    IN: old accelerations, OUT: new accelerations
 *
 * The "update position" and "update velocity" execute 12 floating point
 * operations per iteration each, resulting in 24*n floating point ops
 * together.  Add to it the 3 flops to precompute time step related values and
 * the 13*n*(n-1) flops from acceleration() to arrive at the total of
 * 13*n*(n-1)+24*n+2 floating point operations per call to leapfrog().
 */
void leapfrog (const int N, double dt, vector3 r[], vector3 v[], double m[],
               vector3 a[])
{
    int i;
    vector3 aold[N]; // temporary variable
    
    /* precompute some values from the time step */
    double dt2 = dt*dt*0.5;
    double dthalf = dt*0.5;
    
    /* update position */
    for (i=0; i<N; i++)
    {
        r[i][0] += dt*v[i][0] + dt2*a[i][0];
        r[i][1] += dt*v[i][1] + dt2*a[i][1];
        r[i][2] += dt*v[i][2] + dt2*a[i][2];
    }

    /* save and clear acceleration */
    for (i=0; i<N; i++)
    {
        aold[i][0] = a[i][0];
        aold[i][1] = a[i][1];
        aold[i][2] = a[i][2];
        a[i][0] = a[i][1] = a[i][2] = 0.0;
    }

    /* compute new acceleration */
    acceleration(N,r,m,a);

    /* update velocity */
    for (i=0; i<N; i++)
    {
        v[i][0] += dthalf*aold[i][0] + dthalf*a[i][0];
        v[i][1] += dthalf*aold[i][1] + dthalf*a[i][1];
        v[i][2] += dthalf*aold[i][2] + dthalf*a[i][2];
    }
}


/**
 * main program
 */
int main ()
{
    /* number of bodies */
    const int N = 200;
    /* number of time steps to do */
    const int timesteps = 1000;
    /* delta t */
    const double dt = 1E-2;
    /* when to measure MFLOP rate and write output files: every mod time steps */
    const int mod = 10;
    /* mode: leapfrog (1) or euler (2) */
    const int mode = 2;

    /* arrays for position, velocity, mass, and acceleration of the bodies */
    vector3 r[N];
    vector3 v[N];
    double  m[N];
    vector3 a[N];
    
    /* loop counter */
    int i;
    /* count the time steps done */
    int k = 0;
    /* accumulated simulation time */
    double t = 0;
    /* for measuring the elapsed time and calculating the MFLOP rate */
    double start,stop,elapsed,flop;

    /* get initial values from one of the generator functions */
    plummer(N,r,v,m);
    
    /* initialise the acceleration */
    for (i=0; i<N; i++)
        a[i][0] = a[i][1] = a[i][2] = 0.0;
    acceleration(N,r,m,a);

    /* write out the initial values */
    nbody_write(N,r,v,m,t,dt, k/mod);

    /* start the time measurement */
    start = get_time();
    for (k=1; k<timesteps; k++) {
        /* do one timestep and increment elapsed simulated time */
        if (mode == 1)
        {
            leapfrog(N,dt,r,v,m,a);
        }
        else
        {
            for (i=0; i<N; i++)
                a[i][0] = a[i][1] = a[i][2] = 0.0;
            acceleration(N,r,m,a);
            euler(N,dt,r,v,m,a);
        }
        t += dt;
        if (k%mod==0) {
            /* write statistics */
            stop = get_time();
            elapsed = stop-start;
            /* 13*n*(n-1)+24*n+3 FLOP from leaprog(), 1 from the t+=dt above */
            if (mode == 1)
            {
                flop = mod*(13.0*N*(N-1.0)+24.0*N+4.0);
            }
            else
            {
                flop = mod*(13.0*N*(N-1.0)+12.0*N);
            }   
            printf("%g seconds for %g ops = %g MFLOPS\n",
                elapsed,flop,flop/elapsed/1E6);
            
            /* write output file */
            nbody_write(N,r,v,m,t,dt, k/mod);
            
            /* restart the timer */
            start = get_time();
        }
    }

    /* we are done */
    return 0;
}