/*
 * Multiple bodies problem
 */

#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		NUM_OBJECTS		5
ph_obj_t	obj[NUM_OBJECTS + 100];

/*
 * 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;
}


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,1);
	vec3Set(obj[1].vel,-1,0,0);

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

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

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


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

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

#if 0

/*
 * Dumb implementation of Euler method on non-adaptive time-domain grid.
 * Not working, of course :)
 */

void euler_step(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);
		vec3Add(obj[i].vel,obj[i].vel,v);
		vec3Scale(v,obj[i].vel,timediff);
		vec3Add(obj[i].pos,obj[i].pos,v);
	}
}

#endif

/*
 * Runge corrector (refining time-domain grid). Works well, but:
 * TODO: implement real Richardson extrapolation.
 */

void solver_retrieve_phase(vec3_t dest_pos[],vec3_t dest_vel[])
{
	int i;
	for(i=0;i<NUM_OBJECTS;i++)
	{
		vec3Copy(dest_pos[i],obj[i].pos);
		vec3Copy(dest_vel[i],obj[i].vel);
	}
}

void solver_store_phase(vec3_t src_pos[],vec3_t src_vel[])
{
	int i;
	for(i=0;i<NUM_OBJECTS;i++)
	{
		vec3Copy(obj[i].pos,src_pos[i]);
		vec3Copy(obj[i].vel,src_vel[i]);
	}
}

void solver_copy_phase(vec3_t dest_pos[],vec3_t dest_vel[],
	vec3_t src_pos[],vec3_t src_vel[])
{
	int i;
	for(i=0;i<NUM_OBJECTS;i++)
	{
		vec3Copy(dest_vel[i],src_vel[i]);
		vec3Copy(dest_pos[i],src_pos[i]);
	}
}

void euler_update_phase(vec3_t pos[],vec3_t vel[],vec3_t frc[],vec_t timediff)
{
	int i;
	vec3_t v;
	for(i=0;i<NUM_OBJECTS;i++)
	{
		vec3Scale(v,frc[i],obj[i].rm * timediff);
		vec3Add(vel[i],vel[i],v);
		vec3Scale(v,vel[i],timediff);
		vec3Add(pos[i],pos[i],v);
	}
}

void lerp_vectors(vec3_t dest_lerp[],vec3_t lerp_with[],vec_t midpoint)
{
	int i;
	vec3_t v;
	for(i=0;i<NUM_OBJECTS;i++)
	{
		vec3Sub(v,lerp_with[i],dest_lerp[i]);
		vec3Scale(v,v,midpoint);
		vec3Add(dest_lerp[i],dest_lerp[i],v);
	}
}

vec3_t	startpos_forces[NUM_OBJECTS];
vec3_t	midpos_forces[NUM_OBJECTS];

vec3_t	mid_pos[NUM_OBJECTS];
vec3_t	mid_vel[NUM_OBJECTS];

vec3_t	end_pos[NUM_OBJECTS];
vec3_t	end_vel[NUM_OBJECTS];

vec3_t	tmp_pos[NUM_OBJECTS];
vec3_t	tmp_vel[NUM_OBJECTS];

vec_t simulation_time = 0;
u32_t cycles = 0;

vec_t phasespace_distance(vec3_t pos1[],vec3_t vel1[],vec3_t pos2[],vec3_t vel2[])
{
	int i;
	vec_t sum = 0;
	vec3_t diff;
	for(i=0;i<NUM_OBJECTS;i++)
	{
		vec3Sub(diff,pos1[i],pos2[i]);
		sum += vec3Dot(diff,diff);
		vec3Sub(diff,vel1[i],vel2[i]);
		sum += vec3Dot(diff,diff);
	}
	sum = sqrt(sum);
	return sum;
}

vboolean step_too_large(vec3_t pos1[],vec3_t vel1[],vec3_t pos2[],vec3_t vel2[])
{
	vec_t dist = phasespace_distance(pos1,vel1,pos2,vel2);
	if(dist > 1e-6)
		return true;
	return false;
}

vec_t last_time_diff;

void solver_step(vec_t time_diff)
{
	vec_t half_diff;

	solver_retrieve_phase(mid_pos,mid_vel);
	solver_copy_phase(end_pos,end_vel,mid_pos,mid_vel);

	compute_forces_nip(mid_pos,startpos_forces);
	/* First trial of full step */
	euler_update_phase(end_pos,end_vel,startpos_forces,time_diff);

	for(;;)
	{
		cycles++;
		half_diff = time_diff * 0.5;

		euler_update_phase(mid_pos,mid_vel,startpos_forces,half_diff);
		solver_copy_phase(tmp_pos,tmp_vel,mid_pos,mid_vel);
		compute_forces_nip(mid_pos,midpos_forces);
		euler_update_phase(mid_pos,mid_vel,midpos_forces,half_diff);

		if(!step_too_large(end_pos,end_vel,mid_pos,mid_vel))
			break;

		solver_retrieve_phase(mid_pos,mid_vel);
		solver_copy_phase(end_pos,end_vel,tmp_pos,tmp_vel);

		time_diff = half_diff;
	}
	/* lerp result (this is stupid; see todo) */
#define	LERP_BACKPOS	(vec_t)(-1)
	lerp_vectors(mid_pos,end_pos,LERP_BACKPOS);
	lerp_vectors(mid_vel,end_vel,LERP_BACKPOS);

	solver_store_phase(mid_pos,mid_vel);

	simulation_time += time_diff;
	last_time_diff = time_diff;
}

vec_t initial_energy;

void solver_init()
{
	initial_value();
	initial_energy = full_energy();
	last_time_diff = 0.0005f;
}

void solver_solve()
{
	int i;
	for(i=0;i<10;i++)
	{
		solver_step(last_time_diff);
		/* attempt to speed up calculation */
		if(last_time_diff < 0.0001)
			last_time_diff *= 2;
	}
}

vec_t max_diff = 0;
vec_t energy, diff;

void solver_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()
{
	solver_status();
	begin_overlay(40,30);
	vglPrintf("Cycle: %u\nTime: %4.3g\nFull energy: %4.5g\nDifference: %g\n"
		"Max. en. diff.: %g",
		cycles,simulation_time,energy,diff,max_diff);
	end_overlay();
}

GLfloat alpha = 0, betta = 0;

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

	/* 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);
    glRotatef(betta,0,0,1);
	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();
	glLineWidth(1);
	glBegin(GL_LINES);
		for(i=0;i<NUM_OBJECTS;i++)
		{
			VGL_VERTEX3V(obj[i].pos);
			glVertex3f((GLfloat)obj[i].pos[0],(GLfloat)obj[i].pos[1],0);
		}
	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;
				case key_up:
					alpha += 10;
					break;
				case key_down:
					alpha -= 10;
					break;
				case key_left:
					betta += 10;
					break;
				case key_right:
					betta -= 10;
					break;
				case key_space:
					alpha = betta = 0;
					break;
			}
			break;
		default:
			break;
	}
	return false;
}

void main_loop(void)
{
	input_record_t in;
	solver_init();
	for(;;)
	{
		while(vglPending())
		{
			vglNextEvent(&in);
			if(handle_event_check_break(&in))
				goto good_bye;
		}
		solver_solve();
		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);
		solver_init();
		for(st=0;;st++)
		{
			for(i=0;i<10;i++)
				solver_solve();
			solver_status();
			if((timeTickCount() - starttime) > 1000)
			{
				starttime = timeTickCount();
				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 (Runge corrector)",
			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());
	}
	solver_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;

}
