/* A sequential implementation of a of a n-body simulation.
 *
 * Your task is to investigate how to parallelized this program using
 * {OpenMP,MPI,OpenCL} and also report on your experiences.  Feel free
 * to rewrite anything you deem necessary.
 *
 * Report any issues or mistakes in this file to pkal@cs.ru.nl. */

#define _POSIX_SOURCE		/* for strtok_r */
#include <getopt.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#include <assert.h>
#include <stdbool.h>

/* Compile-time configuration options: */

#ifdef DIMENSIONS
#define N DIMENSIONS		/* number of dimensions */
#else
#define N 3			/* ... default to 3 */
#endif
static_assert(1 <= N, "You need at least one dimensions.");

/* The specification may omit the mass of a body.  In that case we
   fall back to just assuming the "unit" mass. */
#ifndef DEFAULT_MASS
#define DEFAULT_MASS 1
#endif


/* "generic" implementation of basic linear algebra operations. */

struct vec { double v[N]; };

static void
vec_mul(struct vec * const vec, const double scalar)
{
  for (int i = 0; i < N; ++i) {
    vec->v[i] *= scalar;
  }
}

static void
vec_add(struct vec * const restrict acc,
	const struct vec * const restrict vec)
{
  for (int i = 0; i < N; ++i) {
    acc->v[i] += vec->v[i];
  }
}

static double
vec_mag(const struct vec * const restrict v)
{
  double sum = 0;
  for (int i = 0; i < N; ++i) {
    sum += pow(v->v[i], 2);
  }
  return sqrt(sum);
}

static void
vec_copy(const struct vec * const restrict src,
	 struct vec * const restrict dst)
{
  memcpy(dst, src, sizeof *dst);
}

static void
vec_read(char *str, struct vec * const restrict v)
{
  char *endp;
  for (int i = 0; i < N; ++i) {
    char *tok = strtok_r(str, ", ", &endp); str = NULL;
    if (tok != NULL) { sscanf(tok, "%lf", &(v->v[i])); }
  }
}

static void
vec_write(FILE *f, const struct vec * const restrict v)
{
  for (int i = 0; i < N; ++i) {
    assert(!isnan(v->v[i]));
    fprintf(f, "%s%g", i > 0 ? " " : "", v->v[i]);
  }
  fprintf(f, "\n");
}

#define ZERO (&((struct vec) { 0 }))


/* The actual simulation */

struct body {
  double mass;
  struct vec pos;		/* coordinate */
  struct vec vel;		/* velocity */
  struct vec acc[2];		/* acceleration (interm storage) */
};

struct conf {
  unsigned count;		/* number of bodies */
  struct body *bodies;		/* array of body objects */
  double dt;			/* time step */
  double G; 			/* gravitational constant */
};

/* Determine the gravitational forces for each body, remembering the
   the force in the previous iteration.  Recall that classically, the
   gravitiational force is determined directly by the positions of all
   other bodies in the simulation, and not their velocity. */
static void
accelerate(const struct conf * const restrict conf)
{
  /* aᵢ = G ∑ⱼ(i≠j), mⱼ ((xⱼ - xᵢ) / (xⱼ - xᵢ)³) */
  for (unsigned i = 0; i < conf->count; ++i) {
    vec_copy(&(conf->bodies[i].acc[0]), &(conf->bodies[i].acc[1]));
    vec_copy(ZERO, &(conf->bodies[i].acc[0]));
    for (unsigned j = 0; j < conf->count; ++j) {
      if (i == j) continue;

      struct vec tmp;
      vec_copy(&(conf->bodies[i].pos), &tmp);
      vec_mul(&tmp, -1);
      vec_add(&tmp, &(conf->bodies[j].pos));

      double dist = vec_mag(&tmp);
      vec_mul(&tmp, conf->bodies[j].mass / pow(dist, 3));
      vec_add(&(conf->bodies[i].acc[0]), &tmp);
    }

    vec_mul(&(conf->bodies[i].acc[0]), conf->G);
  }
}

/* Integration-step updating the position and velocity of each body
   depending on their current position and velocity.  We are using the
   "Leapfrog" method that depends on the current and the previous
   acceleration:

     xᵢ₊₁ = xᵢ + vᵢ Δt + ½aᵢ Δt²
     vᵢ₊₁ = vᵢ + ½(aᵢ + aᵢ₊₁) Δt

   See https://en.wikipedia.org/wiki/Leapfrog_integration */
static void
step(const struct conf * const conf)
{
  struct vec tmp;

  accelerate(conf);
  for (unsigned i = 0; i < conf->count; ++i) {
    vec_copy(&(conf->bodies[i].vel), &tmp);
    vec_mul(&tmp, conf->dt);
    vec_add(&(conf->bodies[i].pos), &tmp);

    vec_copy(&(conf->bodies[i].acc[0]), &tmp);
    vec_mul(&tmp, .5 * pow(conf->dt, 2));
    vec_add(&(conf->bodies[i].pos), &tmp);
  }

  accelerate(conf);
  for (unsigned i = 0; i < conf->count; ++i) {
    vec_copy(&(conf->bodies[i].acc[0]), &tmp);
    vec_add(&tmp, &(conf->bodies[i].acc[1]));
    vec_mul(&tmp, .5 * conf->dt);
    vec_add(&(conf->bodies[i].vel), &tmp);
  }
}


/* The remaining file takes care of administrative business. */

/* ASCII to positive integer (or die) */
static unsigned long
atop(char *s)
{
  char *endp;

  errno = 0;
  unsigned long n = strtoul(s, &endp, 10);
  if (errno != 0) {
    perror("strtoul");
    exit(EXIT_FAILURE);
  }
  if (*endp != '\0') {
    fprintf(stderr, "Invalid number \"%s\"\n", s);
    exit(EXIT_FAILURE);
  }

  return n;
}

#define USAGE							\
  "Usage %s [options...] [bodies...], where options are:\n"	\
  " -t TIME   Set the size of Δt.\n"				\
  " -n COUNT  Number of Δt timesteps to simulate.\n"		\
  " -s COUNT  How infrequently to dump the state.\n"		\
  " -G REAL   Custom gravitational constant.\n"			\
  " -x        Non-stop mode, ignore -n.\n"			\
  " -v        Print progress.\n"				\
  " -h        Print this overview.\n\n"				\
  "The specification of a body has the form:\n\n"		\
  "    x1,x2,x3,...;v1,v2,v3,...;m\n\n"				\
  "where x1,x2,x3,... is the initial position, and\n"		\
  "v2,v3,v3,... is the initial velocity, m is the mass.\n"	\

int
main(int argc, char **argv)
{
  struct conf conf = {
    .dt = 1e-3,
    .G  = 6.67430e-11, /* https://physics.nist.gov/cgi-bin/cuu/Value?bg */
  };
  unsigned long steps = 1000000000, sample = 100000;
  int status = EXIT_SUCCESS;
  bool nonstop = false, verbose = false;

  int opt;
  while ((opt = getopt(argc, argv, "t:n:s:G:xvh")) != -1) {
    switch (opt) {
    case 't': conf.dt = atof(optarg); break;
    case 'n': steps = atop(optarg); break;
    case 's': sample = atop(optarg); break;
    case 'G': conf.G = atof(optarg); break;
    case 'x': nonstop = true; break;
    case 'v': verbose = true; break;
    case 'h': goto usage;
    default: exit(EXIT_FAILURE);
    }
  }

  /* the remaining arguments are body specifications */
  if ((conf.count = argc - optind) == 0) {
    status = EXIT_FAILURE;
    goto usage;
  }

  conf.bodies = calloc(conf.count, sizeof (struct body));
  if (NULL == conf.bodies) {
    perror("calloc");
    exit(EXIT_FAILURE);
  }
  for (unsigned i = 0; i < conf.count; ++i) {
    conf.bodies[i] = (struct body) { /* default initialisation */
      .mass   = DEFAULT_MASS,
    };

    char *tok, *s;
    if ((tok = strtok_r(argv[optind + i], ";\t", &s)) == NULL) continue;
    vec_read(tok, &(conf.bodies[i].pos));

    if ((tok = strtok_r(NULL, ";\t", &s)) == NULL) continue;
    vec_read(tok, &(conf.bodies[i].vel));

    if ((tok = strtok_r(NULL, ";\t", &s)) == NULL) continue;
    sscanf(tok, "%lf", &(conf.bodies[i].mass));
  }

  double t = 0;
  for (unsigned long i = 0; nonstop || i < steps; ++i, t += conf.dt) {
    if (i % sample == 0) {
      for (unsigned j = 0; j < conf.count; ++j) {
	vec_write(stdout, &(conf.bodies[j].pos));
      }
      printf("\n\n");		/* Gnuplot block seperator */

      if (verbose) {
	fprintf(stderr, "\r%3.f%%, t = %f",  floor((double)i / steps * 1e2), t);
      }
    }
    step(&conf);
  }
  if (verbose) {
    fprintf(stderr, "100%%, t = %f\n", t);
  }

  free(conf.bodies);
  return status;

 usage:
  fprintf(stderr, USAGE , argv[0]);
  return status;
}
