/*******************************************************

* Copyright (c) 2014, ArrayFire

* All rights reserved.

*

* This file is distributed under 3-clause BSD license.

* The complete license agreement can be obtained at:

* http://arrayfire.com/licenses/BSD-3-Clause

********************************************************/

#include <arrayfire.h>

#include <cstdio>

#include <iostream>

#include <vector>

#include "gravity_sim_init.h"

using namespace af;

using namespace std;

static const bool is3D = true;

const static int total_particles = 4000;

static const int reset = 3000;

static const float min_dist = 3;

static const int width = 768, height = 768, depth = 768;

static const int gravity_constant = 20000;

float mass_range = 0;

float min_mass = 0;

void initial_conditions_rand(af::array &mass, vector<af::array> &pos,

vector<af::array> &vels,

vector<af::array> &forces) {

for (int i = 0; i < (int)pos.size(); ++i) {

pos[i] = af::randn(total_particles) * width + width;

vels[i] = 0 * af::randu(total_particles) - 0.5;

forces[i] = af::constant(0, total_particles);

}

mass = af::constant(gravity_constant, total_particles);

}

void initial_conditions_galaxy(af::array &mass, vector<af::array> &pos,

vector<af::array> &vels,

vector<af::array> &forces) {

af::array initial_cond_consts(af::dim4(7, total_particles), hbd);

initial_cond_consts = initial_cond_consts.T();

for (int i = 0; i < (int)pos.size(); ++i) {

pos[i] = af::randn(total_particles) * width + width;

vels[i] = 0 * (af::randu(total_particles) - 0.5);

forces[i] = af::constant(0, total_particles);

}

mass = initial_cond_consts(span, 0);

pos[0] = (initial_cond_consts(span, 1) / 32 + 0.6) * width;

pos[1] = (initial_cond_consts(span, 2) / 32 + 0.3) * height;

pos[2] = (initial_cond_consts(span, 3) / 32 + 0.5) * depth;

vels[0] = (initial_cond_consts(span, 4) / 32) * width;

vels[1] = (initial_cond_consts(span, 5) / 32) * height;

vels[2] = (initial_cond_consts(span, 6) / 32) * depth;

pos[0](seq(0, pos[0].dims(0) - 1, 2)) -= 0.4 * width;

pos[1](seq(0, pos[0].dims(0) - 1, 2)) += 0.4 * height;

vels[0](seq(0, pos[0].dims(0) - 1, 2)) += 4;

min_mass = min<float>(mass);

mass_range = max<float>(mass) - min<float>(mass);

}

af::array ids_from_pos(vector<af::array> &pos) {

return (pos[0].as(u32) * height) + pos[1].as(u32);

}

af::array ids_from_3D(vector<af::array> &pos, float Rx, float Ry, float Rz) {

af::array x0 = (pos[0] - width / 2);

af::array y0 =

(pos[1] - height / 2) * cos(Rx) + (pos[2] - depth / 2) * sin(Rx);

af::array z0 =

(pos[2] - depth / 2) * cos(Rx) - (pos[2] - depth / 2) * sin(Rx);

af::array x1 = x0 * cos(Ry) - z0 * sin(Ry);

af::array y1 = y0;

af::array x2 = x1 * cos(Rz) + y1 * sin(Rz);

af::array y2 = y1 * cos(Rz) - x1 * sin(Rz);

x2 += width / 2;

y2 += height / 2;

return (x2.as(u32) * height) + y2.as(u32);

}

af::array ids_from_3D(vector<af::array> &pos, float Rx, float Ry, float Rz,

af::array filter) {

af::array x0 = (pos[0](filter) - width / 2);

af::array y0 = (pos[1](filter) - height / 2) * cos(Rx) +

(pos[2](filter) - depth / 2) * sin(Rx);

af::array z0 = (pos[2](filter) - depth / 2) * cos(Rx) -

(pos[2](filter) - depth / 2) * sin(Rx);

af::array x1 = x0 * cos(Ry) - z0 * sin(Ry);

af::array y1 = y0;

af::array x2 = x1 * cos(Rz) + y1 * sin(Rz);

af::array y2 = y1 * cos(Rz) - x1 * sin(Rz);

x2 += width / 2;

y2 += height / 2;

return (x2.as(u32) * height) + y2.as(u32);

}

void simulate(af::array &mass, vector<af::array> &pos, vector<af::array> &vels,

vector<af::array> &forces, float dt) {

for (int i = 0; i < (int)pos.size(); ++i) {

pos[i] += vels[i] * dt;

pos[i].eval();

}

// calculate forces to each particle

vector<af::array> diff(pos.size());

af::array dist = af::constant(0, pos[0].dims(0), pos[0].dims(0));

for (int i = 0; i < (int)pos.size(); ++i) {

diff[i] = tile(pos[i], 1, pos[i].dims(0)) -

transpose(tile(pos[i], 1, pos[i].dims(0)));

dist += (diff[i] * diff[i]);

}

dist = sqrt(dist);

dist = af::max(min_dist, dist);

dist *= dist * dist;

for (int i = 0; i < (int)pos.size(); ++i) {

// calculate force vectors

forces[i] = diff[i] / dist;

forces[i].eval();

// af::array idx = af::where(af::isNaN(forces[i]));

// if(idx.elements() > 0)

// forces[i](idx) = 0;

// forces[i] = sum(forces[i]).T();

forces[i] = matmul(forces[i].T(), mass);

// update force scaled to time, magnitude constant

forces[i] *= (gravity_constant);

forces[i].eval();

// update velocities from forces

vels[i] += forces[i] * dt;

vels[i].eval();

// noise

// forces[i] += 0.1 * af::randn(forces[i].dims(0));

// dampening

// vels[i] *= 1 - (0.005*dt);

}

}

void collisions(vector<af::array> &pos, vector<af::array> &vels, bool is3D) {

// clamp particles inside screen border

af::array invalid_x = -2 * (pos[0] > width - 1 || pos[0] < 0) + 1;

af::array invalid_y = -2 * (pos[1] > height - 1 || pos[1] < 0) + 1;

// af::array invalid_x = (pos[0] < width-1 || pos[0] > 0);

// af::array invalid_y = (pos[1] < height-1 || pos[1] > 0);

vels[0] = invalid_x * vels[0];

vels[1] = invalid_y * vels[1];

af::array projected_px = min(width - 1, max(0, pos[0]));

af::array projected_py = min(height - 1, max(0, pos[1]));

pos[0] = projected_px;

pos[1] = projected_py;

if (is3D) {

af::array invalid_z = -2 * (pos[2] > depth - 1 || pos[2] < 0) + 1;

vels[2] = invalid_z * vels[2];

af::array projected_pz = min(depth - 1, max(0, pos[2]));

pos[2] = projected_pz;

}

}

int main(int, char **) {

try {

af::info();

af::Window myWindow(width, height,

"Gravity Simulation using ArrayFire");

myWindow.setColorMap(AF_COLORMAP_HEAT);

int frame_count = 0;

// Initialize the kernel array just once

const af::array draw_kernel = gaussianKernel(7, 7);

const int dims = (is3D) ? 3 : 2;

vector<af::array> pos(dims);

vector<af::array> vels(dims);

vector<af::array> forces(dims);

af::array mass;

// Generate a random starting state

initial_conditions_galaxy(mass, pos, vels, forces);

af::array image = af::constant(0, width, height);

af::array ids(total_particles, u32);

af::timer timer = af::timer::start();

while (!myWindow.close()) {

float dt = af::timer::stop(timer);

timer = af::timer::start();

af::array mid = mass(span) > (min_mass + mass_range / 3);

ids = (is3D) ? ids_from_3D(pos, 0, 0 + frame_count / 150.f, 0, mid)

: ids_from_pos(pos);

// ids = (is3D)? ids_from_3D(pos, 0, 0, 0, mid) : ids_from_pos(pos);

// //uncomment for no 3d rotation

image(ids) += 4.f;

mid = mass(span) > (min_mass + 2 * mass_range / 3);

ids = (is3D) ? ids_from_3D(pos, 0, 0 + frame_count / 150.f, 0, mid)

: ids_from_pos(pos);

// ids = (is3D)? ids_from_3D(pos, 0, 0, 0, mid) : ids_from_pos(pos);

// //uncomment for no 3d rotation

image(ids) += 4.f;

ids = (is3D) ? ids_from_3D(pos, 0, 0 + frame_count / 150.f, 0)

: ids_from_pos(pos);

// ids = (is3D)? ids_from_3D(pos, 0, 0, 0) : ids_from_pos(pos);

// //uncomment for no 3d rotation

image(ids) += 4.f;

image = convolve(image, draw_kernel);

myWindow.image(image);

image = af::constant(0, image.dims());

frame_count++;

// Generate a random starting state

if (frame_count % reset == 0) {

initial_conditions_galaxy(mass, pos, vels, forces);

}

// simulate

simulate(mass, pos, vels, forces, dt);

// check for collisions and adjust positions/velocities accordingly

collisions(pos, vels, is3D);

}

} catch (af::exception &e) {

fprintf(stderr, "%s\n", e.what());

throw;

}

return 0;

}