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

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

/*
 * Global constants
 */

/* physical parameters */
const double mass_sun = 1.988435e30;
const double pc_in_m = 3.08567758129e16;
const double gamma_si = 6.67428e-11;
/* set gravity constant */
const double G = gamma_si/(pc_in_m*pc_in_m*pc_in_m)*mass_sun*(365.25*86400)*(365.25*86400);

/* The epsilon^2 for the softening */
const double epsilon2 = 1E-3;

/* Block size for tiling */
const int B = 50;

/**
 * Compute acceleration (with tiling)
 *
 * \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,i,j,iend,jstart,jend;
    double d0,d1,d2,d,d_sq,factori,factorj;
    double invfact;
    vector3 Ri[B], Rj[B], Ai[B], Aj[B];
    double  Mi[B], Mj[B];
    
    /* compute acceleration exploiting symmetry */
    for (I=0; I<n; I+=B) {
        for (J=I; J<n; J+=B) {
            /* store local copy of the blocks to avoid aliasing */
            iend = std::min(I+B,n);
            jend = std::min(J+B,n);
            for (i=I; i<iend; i++) {
                Ri[i-I][0] = r[i][0];
                Ri[i-I][1] = r[i][1];
                Ri[i-I][2] = r[i][2];
                Ai[i-I][0] = 0.0;
                Ai[i-I][1] = 0.0;
                Ai[i-I][2] = 0.0;
                Mi[i-I] = m[i];
            }
            for (j=J; j<jend; j++) {
                Rj[j-J][0] = r[j][0];
                Rj[j-J][1] = r[j][1];
                Rj[j-J][2] = r[j][2];
                Aj[j-J][0] = 0.0;
                Aj[j-J][1] = 0.0;
                Aj[j-J][2] = 0.0;
                Mj[j-J] = m[j];
            }
            /* compute locally */
            for (i=0; i<iend-I; i++) {
                jstart = std::max(0,I+i+1-J);
                for (j=jstart; j<jend-J; j++) {
                    /* compute distance vector */
                    d0 = Rj[j][0]-Ri[i][0];
                    d1 = Rj[j][1]-Ri[i][1];
                    d2 = Rj[j][2]-Ri[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 = Mi[i]*invfact;
                    factorj = Mj[j]*invfact;
            
                    /* compute acceleration of body i due to body j */
                    Ai[i][0] += factorj*d0;
                    Ai[i][1] += factorj*d1;
                    Ai[i][2] += factorj*d2;
            
                    /* compute acceleration of body j due to body i */
                    Aj[j][0] -= factori*d0;
                    Aj[j][1] -= factori*d1;
                    Aj[j][2] -= factori*d2;
                }
            }
            /* store result */
            for (i=I; i<iend; i++) {
                a[i][0] += Ai[i-I][0];
                a[i][1] += Ai[i-I][1];
                a[i][2] += Ai[i-I][2];
            }
            for (j=J; j<jend; j++) {
                a[j][0] += Aj[j-J][0];
                a[j][1] += Aj[j-J][1];
                a[j][2] += Aj[j-J][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]     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 = new vector3[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];
    }

    delete[] aold;
}


/**
 * main program
 */
int main (int argc, char** argv)
{
    /* number of bodies */
    int N = 200;
    /* delta t */
    double dt = 1E-2;
    /* when to measure MFLOP rate and write output files: every mod time steps */
    int mod = 10;
    /* mode: leapfrog (1) or euler (2) */
    const int mode = 1;
    /* measure time */
    double T, Tend;
    /* number of parallel processes */
    int numthreads = 1;

    /* read commandline parameters */
    if (argc!=6)
	{
        printf("usage: %s <nbodies> <Tend> <every> <dt> <numthreads>\n", argv[0]);
        return 1;
	}
    sscanf(argv[1],"%d",&N);
    sscanf(argv[2],"%lf",&Tend);
    sscanf(argv[3],"%d",&mod);
    sscanf(argv[4],"%lf",&dt);
    sscanf(argv[5],"%d",&numthreads);

    /* arrays for position, velocity, mass, and acceleration of the bodies */
    vector3* r = new vector3[N];
    vector3* v = new vector3[N];
    double*  m = new double[N];
    vector3* a = new vector3[N];
    
    /* loop counter */
    int i;
    /* count the time steps done */
    int k = 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 */
    int seed = 42;
    collision(N,seed,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,0,dt, k/mod);

    /* start the time measurement */
    start = get_time();
    for (T=0, k=1; T<Tend; T+=dt, 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);
        }
        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,dt, k/mod);
            
            /* restart the timer */
            start = get_time();
        }
    }

    delete[] r;
    delete[] v;
    delete[] m;
    delete[] a;

    /* we are done */
    return 0;
}