#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <pthread.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 = 90;

typedef struct acceleration_data {
    int n,rank,numthreads;
    double *m;
    vector3 *a, *r;
} acceleration_data;

pthread_mutex_t *mutex_a;

/**
 * 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 (acceleration_data *data)
{
    int n, rank, numthreads;
    vector3 *r, *a;
    double *m;
    
    int I,J,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];
    
    /* extract data */
    n = data->n;    
    rank = data->rank;
    numthreads = data->numthreads;
    r = data->r;
    a = data->a;
    m = data->m;
    
    /* compute acceleration exploiting symmetry */
    for (I_=0; I_<=n/B; I_++) {
        for (J_=I_; J_<=n/B; J_++) {
            if ((I_+J_) % numthreads == rank) {
                /*printf("i am thread %d of %d and i compute block (%d,%d)\n",rank,numthreads,I_,J_);*/
                I=I_*B;
                J=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++) {
                    pthread_mutex_lock(&mutex_a[i]);
                    a[i][0] += Ai[i-I][0];
                    a[i][1] += Ai[i-I][1];
                    a[i][2] += Ai[i-I][2];
                    pthread_mutex_unlock(&mutex_a[i]);
                }
                for (j=J; j<jend; j++) {
                    pthread_mutex_lock(&mutex_a[j]);
                    a[j][0] += Aj[j-J][0];
                    a[j][1] += Aj[j-J][1];
                    a[j][2] += Aj[j-J][2];
                    pthread_mutex_unlock(&mutex_a[j]);
                }
            }
        }
    }
}


void acceleration_par(int numthreads, int n, vector3 *r, double *m, vector3 *a) {
	pthread_t threads[numthreads];
	acceleration_data data[numthreads];
	int i;
    
    /* initialize data-structs */
    for (i=0; i<numthreads; i++) {
        data[i].n = n;
        data[i].rank = i;
        data[i].numthreads = numthreads;
        data[i].m = m;
        data[i].a = a;
        data[i].r = r;
    }
    
    /* create threads */
    for (i=0; i<numthreads; i++) {
        /*printf("starting thread %d\n",i);*/
        pthread_create(threads+i,NULL,(void*(*)(void*)) acceleration,(void *)(data+i));
    }
    
    /* wait for threads to finish */
    for (i=0; i<numthreads; i++) {
        pthread_join(threads[i],NULL); 
        /* printf("thread %d finished\n",i); */
    }
}

/**
 * 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 numthreads,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_par(numthreads,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);
    
    mutex_a = new pthread_mutex_t[N];
    
    /* initialize mutexes */
	for (i=0; i<N; i++)
	    pthread_mutex_init(&mutex_a[i],NULL);
    
    /* initialise the acceleration */
    for (i=0; i<N; i++)
        a[i][0] = a[i][1] = a[i][2] = 0.0;
    acceleration_par(numthreads,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(numthreads,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_par(numthreads,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();
        }
    }
    
    /* destroy mutexes */
	for (i=0; i<N; i++)
	    pthread_mutex_destroy(&mutex_a[i]);
	
	delete[] mutex_a;

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

    /* we are done */
    return 0;
}