/*
 * Multiple bodies problem: Stoermer's rule solver
 */

#include <v_lib.h>

/*
 * -----------------------------------------------------------------------
 * Part 1: physical model
 */

typedef struct ph_obj_s
{
	vec3_t	pos;
	vec3_t	vel;
	vec3_t	force; /* force acting to this object */
	vec_t	mass, rm; /* gravital mass and reciprocal inertial mass */
}	ph_obj_t;

#define	STEP 0.0001

#define		NUM_OBJECTS		5
ph_obj_t	obj[NUM_OBJECTS + 3];

/*
 * F-vect = radius-vect * m1 * m2 / abs(radius-vect) ^ 3
 */
void gravity_force_nip(int subj_ind,int obj_ind,vec3_t pos[],vec3_t force_to_subj)
{
	vec3_t dr;
	vec_t r, rr, f;
	vec3Sub(dr,pos[obj_ind],pos[subj_ind]);
	r = vec3Abs(dr);
	rr = 1/r;
	vec3Scale(dr,dr,rr);
	f = obj[subj_ind].mass * obj[obj_ind].mass / (r * r);
	vec3Scale(force_to_subj,dr,f);
}

void gravity_force(ph_obj_t *subj,ph_obj_t *obj,vec3_t force_to_subj)
{
	vec3_t dr;
	vec_t r, rr, f;
	vec3Sub(dr,obj->pos,subj->pos);
	r = vec3Abs(dr);
	rr = 1/r;
	vec3Scale(dr,dr,rr);
	f = subj->mass * obj->mass / (r * r);
	vec3Scale(force_to_subj,dr,f);
}

/*
 * For each pairs of objects, compute force and accumulate it
 */
void compute_forces(void)
{
	vec3_t force;
	int i,j,lastobj;
	for(i=0;i<NUM_OBJECTS;i++)
		vec3Zero(obj[i].force);
	lastobj = NUM_OBJECTS-1;
	for(i=0;i<lastobj;i++)
	{
		for(j=i+1;j<NUM_OBJECTS;j++)
		{
			gravity_force(&obj[i],&obj[j],force);
			vec3Add(obj[i].force,obj[i].force,force);
			vec3Sub(obj[j].force,obj[j].force,force);
		}
	}
}

void compute_forces_nip(vec3_t positions[],vec3_t forces[])
{
	vec3_t force;
	int i,j,lastobj;
	for(i=0;i<NUM_OBJECTS;i++)
		vec3Zero(forces[i]);
	lastobj = NUM_OBJECTS-1;
	for(i=0;i<lastobj;i++)
	{
		for(j=i+1;j<NUM_OBJECTS;j++)
		{
			gravity_force_nip(i,j,positions,force);
			vec3Add(forces[i],forces[i],force);
			vec3Sub(forces[j],forces[j],force);
		}
	}
}

/*
 * -----------------------------------------------------------------------
 * Part 2: intial conditons
 */

/*
 * Computes full energy of system
 */
vec_t full_energy(void)
{
	int i, j;
	vec_t vel_squared, r, potential, energy = 0;
	vec3_t r_ji;
	/* Kinetic energy: summ_i(m_i * v_i ^ 2) / 2 */
	for(i=0;i<NUM_OBJECTS;i++)
	{
		vel_squared = vec3Dot(obj[i].vel,obj[i].vel);
		energy += obj[i].mass * vel_squared;
	}
	/* Potential energy */
	for(i=0;i<NUM_OBJECTS;i++)
	{
		potential = 0;
		for(j=0;j<NUM_OBJECTS;j++)
		{
			if(i==j)
				continue;
			vec3Sub(r_ji,obj[j].pos,obj[i].pos);
			r = vec3Abs(r_ji);
			potential += obj[j].mass / r;
		}
		energy -= potential * obj[i].mass;
	}
	energy *= (vec_t)0.5;
	return energy;
}

vec_t	initial_energy;

void initial_value(void)
{
	int i;

    /* initial values for masses, positions and velocities */
	obj[0].mass = 5;
	vec3Set(obj[0].pos,0,0,0);
	vec3Set(obj[0].vel,0,0,0);

	obj[1].mass = 1;
	vec3Set(obj[1].pos,1,1,0);
	vec3Set(obj[1].vel,0,-1,0);

	obj[2].mass = 1;
	vec3Set(obj[2].pos,-1,1,0);
	vec3Set(obj[2].vel,1,0,0);

	obj[3].mass = 1;
	vec3Set(obj[3].pos,-1,-1,0);
	vec3Set(obj[3].vel,0,1,0);

	obj[4].mass = 1;
	vec3Set(obj[4].pos,1,-1,0);
	vec3Set(obj[4].vel,-1,0,0);

	/* setup reciprocal inertial masses */
    for(i=0;i<NUM_OBJECTS;i++)
		obj[i].rm = 1.0f / obj[i].mass;
	initial_energy = full_energy();
}

/*
 * -----------------------------------------------------------------------
 * Part 3: DE solver
 */

/*
 * Dumb implementation of Stoermer's rule on non-adaptive time-domain grid.
 */

vec_t	simulation_time = 0;

vec3_t	prev_r[100];

void stoer_first(vec_t timediff)
{
	int i;
	vec3_t v;
	compute_forces();
	for(i=0;i<NUM_OBJECTS;i++)
	{
		vec3Scale(v,obj[i].force,obj[i].rm * timediff / 2);
		vec3Add(obj[i].vel,obj[i].vel,v);
		vec3Scale(v,obj[i].vel,timediff);
		vec3Copy(prev_r[i],obj[i].pos);
		vec3Add(obj[i].pos,obj[i].pos,v);
	}
	simulation_time += timediff;
}

void stoer_next(vec_t timediff)
{
	int i;
	vec3_t v, v1;
	compute_forces();
   	for(i=0;i<NUM_OBJECTS;i++)
   	{
		/* note: we need velocities for consistency check only ! */
		vec3Sub(v,obj[i].pos,prev_r[i]);
		vec3Scale(v,v,1/timediff);
		vec3Scale(v1,obj[i].force,obj[i].rm * timediff * (vec_t)0.5);
		vec3Add(obj[i].vel,v,v1);

   		vec3Scale(v,obj[i].force,obj[i].rm * timediff * timediff);
   		vec3Scale(v1,obj[i].pos,2);
   		vec3Sub(v1,v1,prev_r[i]);
   		vec3Add(v1,v1,v);
   		vec3Copy(prev_r[i],obj[i].pos);
   		vec3Copy(obj[i].pos,v1);
   	}
   	simulation_time += timediff;
}

vec_t max_diff = 0;
vec_t energy, diff;

void step_status()
{
	energy = full_energy();
	diff = energy - initial_energy;
	if(max_diff < v_abs(diff))
		max_diff = v_abs(diff);
}

/*
 * -----------------------------------------------------------------------
 * Part 4: drawing stuff
 */
void draw_xOy_grid(vec3_t origin,vec_t width,vec_t height,u32_t num_lines)
{
	u32_t i;
	vec3_t p, dp;
	vec3Copy(p,origin);
	p[0] -= (vec_t)0.5 * width;
	p[1] -= (vec_t)0.5 * height;
	vec3Set(dp,width,height,0);
	vec3Scale(dp,dp,(vec_t)1 / (num_lines-1));
	glBegin(GL_LINES);
		/* grid with lines along Oy direction */
		for(i=0;i<num_lines;i++)
		{
			VGL_VERTEX3(p[0]+i*dp[0],p[1],p[2]);
			VGL_VERTEX3(p[0]+i*dp[0],p[1]+height,p[2]);
		}

		/* grid with lines along Ox direction */
		for(i=0;i<num_lines;i++)
		{
			VGL_VERTEX3(p[0],p[1]+i*dp[1],p[2]);
			VGL_VERTEX3(p[0]+width,p[1]+i*dp[1],p[2]);
		}
	glEnd(); /* GL_LINES */
}

void draw_2_xOy_grids(vec3_t origin,vec3_t center,vec_t width,vec_t height,u32_t major_lines,u32_t minor_lines)
{
	vec_t dw, dh;
	vec3_t neworg;
	if(minor_lines < 2 || major_lines < 2)
		return;
	glEnable(GL_BLEND);
	glBlendFunc(GL_SRC_ALPHA,GL_ONE_MINUS_SRC_ALPHA);

	dw = width / (major_lines - 1);
	dh = height / (major_lines - 1);
	vec3Copy(neworg,center);
	neworg[0] -= (vec_t)fmod(center[0] - origin[0],dw);
	neworg[1] -= (vec_t)fmod(center[1] - origin[1],dh);

	glLineWidth(1.5f);
	draw_xOy_grid(neworg,width,height,minor_lines);
	glLineWidth(3.5f);
	draw_xOy_grid(neworg,width,height,major_lines);

	glDisable(GL_BLEND);
}

void begin_overlay(int char_width, int char_height)
{
	glMatrixMode(GL_PROJECTION);
	glPushMatrix();
	glLoadIdentity();
	glOrtho(0,char_width,0,char_height,-1,1);
	glMatrixMode(GL_MODELVIEW);
	glPushMatrix();
	glLoadIdentity();
	VGL_TRANSLATE(0,char_height-1,0);
	glColor4f(1,1,1,1);
	glLineWidth(1);
	glDisable(GL_LINE_SMOOTH);
	glDisable(GL_BLEND);
}

void end_overlay(void)
{
	glMatrixMode(GL_PROJECTION);
	glPopMatrix();
	glMatrixMode(GL_MODELVIEW);
	glPopMatrix();
}

void draw_status()
{
	begin_overlay(30,25);
	step_status();
	vglPrintf("Time: %4.3g\nFull energy: %4.5g\nDifference: %g\n"
		"Max. en. diff.: %g",
		simulation_time,energy,diff,max_diff);
	end_overlay();
}

void draw(void)
{
	int i;
	vec3_t mins, maxs, origin, center;
	vec_t scale;
	glClear(GL_COLOR_BUFFER_BIT);
    glLoadIdentity();

	/* Compute bounding box */
	bboxClear(mins,maxs);
	for(i=0;i<NUM_OBJECTS;i++)
		bboxAddPoint(mins,maxs,obj[i].pos);
	/* translate to center all objects */
	vec3Zero(origin);
	vec3Sub(center,mins,maxs);
	scale = 1.0f / vec3Abs(center);
	if(scale > 1)
		scale = 1;
	vec3Add(center,mins,maxs);
	vec3Scale(center,center,0.5f);
    VGL_SCALE(scale,scale,scale);
	VGL_TRANSLATE(-center[0],-center[1],-center[2]);
	/* Draw coordsys-aligned grid on xOy plane */
	vec3Set(origin,0,0,0);
	glColor4f(.5f,1.0f,.5f,0.7f);
	draw_2_xOy_grids(origin,center,8,8,5,25);
	glColor4f(1,1,1,0.7f);
	/* Origin */
	glBegin(GL_POINTS);
		VGL_VERTEX3V(origin);
	glEnd();
	/* Draw masses */
	glPointSize(6);
	glColor3f(0,1,1);
	glBegin(GL_POINTS);
		for(i=0;i<NUM_OBJECTS;i++)
			VGL_VERTEX3V(obj[i].pos);
	glEnd();
	/* Draw status */
	draw_status();
	vglSwapBuffers();
}

/*
 * -----------------------------------------------------------------------
 * Part 5: UI
 */
vboolean handle_event_check_break(input_record_t *in)
{
	switch(in->event)
	{
		case ev_key_down:
			switch(in->first_code)
			{
				case key_escape:
					return true;
			}
			break;
		default:
			break;
	}
	return false;
}

void main_loop(void)
{
	input_record_t in;
	stoer_first((vec_t)STEP);
	for(;;)
	{
		while(vglPending())
		{
			vglNextEvent(&in);
			if(handle_event_check_break(&in))
				goto good_bye;
		}
		{ int i; for(i=0;i<10;i++) stoer_next((vec_t)STEP); }
		draw();
	}
good_bye:;
}

vboolean nodraw = false;
vboolean noburn = false;
vec_t	diff_en_limit = 1e6;
vec_t	timelimit = 1e6;

int main(int argc,char *argv[])
{
	int i;
	for(i=1;i<argc;i++)
	{
		char *cmd = argv[i];
		if(!strcmp(cmd,"nodraw"))
			nodraw = true;
		else if(!strcmp(cmd,"noburn"))
			noburn = true;
		else if(!strcmp(cmd,"diff_en_limit"))
			diff_en_limit = atof(argv[++i]);
		else if(!strcmp(cmd,"timelimit"))
			timelimit = atof(argv[++i]);
	}
	logStartConsoleLog();
	logUpperLevel(LOG_TRACE);
	initial_value();
	if(nodraw)
	{
		int i,st;
		u32_t starttime = timeTickCount();
		conPrintf("Diff. en. limit %g\nTimelimit %g\n",diff_en_limit,timelimit);
		stoer_first((vec_t)STEP);
		for(st=0;;st++)
		{
			for(i=0;i<10;i++)
				stoer_next((vec_t)STEP);
			if((timeTickCount() - starttime) > 1000)
			{
				starttime = timeTickCount();
				step_status();
				conPrintf("Sim. time: %4.3g\tFull en.: %4.5g\t"
					"Diff.: %g\t\tMax. diff.: %g\n",
					simulation_time,energy,diff,max_diff);
				if(noburn)
					timeDelay(333); /* we use 2/3 of CPU power */
			}
			if(simulation_time > timelimit)
				break;
			if(max_diff > diff_en_limit)
				break;
		}
	}
	else
	{
		vglCreateWindow("Initial value problem for multiple masses",0,0,512,512);
		vglLoadDefaultFont();
		glMatrixMode(GL_PROJECTION);
		glLoadIdentity();
		glOrtho(-1,1,-1,1,-1,1);
		glMatrixMode(GL_MODELVIEW);
		glLoadIdentity();
		glEnable(GL_POINT_SMOOTH);
		glAlphaFunc(GL_GREATER,0.5f);
		glEnable(GL_ALPHA_TEST);
		glEnable(GL_LINE_SMOOTH);
		main_loop();
		vglDestroyWindow(vglGetCurrentWindow());
	}
	step_status();
	conPrintf("Sim. time: %4.3g\tFull en.: %4.5g\t"
		"Diff.: %g\t\tMax. diff.: %g\n",
		simulation_time,energy,diff,max_diff);
	return 0;
}
