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

#define		MAX_DR			10e-8
#define		MAX_DV			10e-9
#define		MAX_TIME_DIFF	0.001
#define		FFORCE			0

void gravity_force(vec3_t force_ret,int index,vec3_t pos,vec3_t other_pos[])
{
	int i;
	vec3_t dr,df;
	vec_t r, rr, f;
	vec3Zero(force_ret);
	for(i=0;i<NUM_OBJECTS;i++)
	{
		if(i == index)
			continue;
		vec3Sub(dr,other_pos[i],pos);
		r = vec3Abs(dr);
		rr = 1 / r;
		vec3Scale(dr,dr,rr);
		f = obj[i].mass * obj[index].mass * rr * rr;
		vec3Scale(df,dr,f);
		vec3Add(force_ret,force_ret,df);
	}
}

/*
 * -----------------------------------------------------------------------
 * 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
 */
void right_side(vec3_t pos_deriv,vec3_t vel_deriv,vec3_t pos,vec3_t all_pos[],vec3_t vel,vec_t fc,int index)
{
	vec3_t tmp;
	vec3_t frc;
	gravity_force(frc,index,pos,all_pos);
	vec3Scale(tmp,vel,fc);
	vec3Sub(tmp,frc,tmp);
	vec3Scale(vel_deriv,tmp,obj[index].rm);
	vec3Copy(pos_deriv,vel);
}

void 
solver_delta(vec3_t delta_p,vec3_t delta_v,vec3_t all_pos[],vec3_t vel,vec_t fc,vec_t timediff,int index)
{
	/*
	 * Euler delta:
	vec3_t pd,vd,pos;
	vec3Copy(pos,all_pos[index]);
	right_side(pd,vd,pos,all_pos,vel,fc,index);
	vec3Scale(delta_v,vd,timediff);
	vec3Scale(delta_p,pd,timediff);
	 */
	/*
	 * RK4 Delta
	 */
	vec3_t dp,dv,pos;
	static vec_t one_by_three = (vec_t)1.0 / (vec_t)3.0;
	vec3_t k1[2], k2[2], k3[2], k4[2], tp, tv;
	vec_t halfstep = 0.5 * timediff;
	vec3Copy(pos,all_pos[index]);
	right_side(k1[0],k1[1],pos,all_pos,vel,fc,index);
	vec3Scale(k1[0],k1[0],halfstep);
	vec3Scale(k1[1],k1[1],halfstep);
	vec3Add(tp,pos,k1[0]);
	vec3Add(tv,vel,k1[1]);
	right_side(k2[0],k2[1],tp,all_pos,tv,fc,index);
	vec3Scale(k2[0],k2[0],halfstep);
	vec3Scale(k2[1],k2[1],halfstep);
	vec3Add(tp,pos,k2[0]);
	vec3Add(tv,vel,k2[1]);
	right_side(k3[0],k3[1],tp,all_pos,tv,fc,index);
	vec3Scale(k3[0],k3[0],timediff);
	vec3Scale(k3[1],k3[1],timediff);
	vec3Add(tp,pos,k3[0]);
	vec3Add(tv,vel,k3[1]);
	right_side(k4[0],k4[1],tp,all_pos,tv,fc,index);
	vec3Scale(k4[0],k4[0],timediff);
	vec3Scale(k4[1],k4[1],timediff);
	vec3Scale(k2[0],k2[0],2);
	vec3Scale(k2[1],k2[1],2);
	vec3Scale(k4[0],k4[0],0.5);
	vec3Scale(k4[1],k4[1],0.5);
	vec3Add(dp,k1[0],k2[0]);
	vec3Add(dp,dp,k3[0]);
	vec3Add(dp,dp,k4[0]);
	vec3Add(dv,k1[1],k2[1]);
	vec3Add(dv,dv,k3[1]);
	vec3Add(dv,dv,k4[1]);
	vec3Scale(delta_p,dp,one_by_three);
	vec3Scale(delta_v,dv,one_by_three);
}

void solver_update_phase(vec3_t pos[],vec3_t vel[],vec_t timediff)
{
	int i;
	vec3_t dv,dp;
	for(i=0;i<NUM_OBJECTS;i++)
	{
		solver_delta(dp,dv,pos,vel[i],FFORCE,timediff,i);
		vec3Add(vel[i],vel[i],dv);
		vec3Add(pos[i],pos[i],dp);
	}
}

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 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	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;

vboolean step_too_large(vec3_t pos1[],vec3_t vel1[],vec3_t pos2[],vec3_t vel2[])
{
	int i;
	vec3_t diff;
	vec_t mr = MAX_DR * MAX_DR;
	vec_t mv = MAX_DV * MAX_DV;
	for(i=0;i<NUM_OBJECTS;i++)
	{
		vec3Sub(diff,pos1[i],pos2[i]);
		if(vec3Dot(diff,diff) > mr)
			return true;
		vec3Sub(diff,vel1[i],vel2[i]);
		if(vec3Dot(diff,diff) > mv)
			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);

	/* First trial of full step */
	solver_update_phase(end_pos,end_vel,time_diff);

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

		solver_update_phase(mid_pos,mid_vel,half_diff);
		solver_copy_phase(tmp_pos,tmp_vel,mid_pos,mid_vel);
		solver_update_phase(mid_pos,mid_vel,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 < MAX_TIME_DIFF)
			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;
	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);

#ifdef AA_ENABLED
	glLineWidth(1.5f);
#endif /* AA_ENABLED */
	draw_xOy_grid(neworg,width,height,minor_lines);
#ifdef AA_ENABLED
	glLineWidth(3.5f);
#endif /* AA_ENABLED */
	draw_xOy_grid(neworg,width,height,major_lines);
}

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);
#ifdef AA_ENABLED
	glLineWidth(1);
#endif /* AA_ENABLED */
}

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();
#ifdef	AA_ENABLED
	glLineWidth(1);
#endif /* AA_ENABLED */
	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)
{
	u32_t time;
	input_record_t in;
	solver_init();
	for(;;)
	{
		while(vglPending())
		{
			vglNextEvent(&in);
			if(handle_event_check_break(&in))
				goto good_bye;
		}
		time = timeTickCount() + 30;
		while(time > timeTickCount())
			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 (RK4 + Runge corrector)",
			0,0,512,512);
		vglLoadDefaultFont();
		glMatrixMode(GL_PROJECTION);
		glLoadIdentity();
		glOrtho(-1,1,-1,1,-1,1);
		glMatrixMode(GL_MODELVIEW);
		glLoadIdentity();
#ifdef	AA_ENABLED
		glEnable(GL_BLEND);
		glBlendFunc(GL_SRC_ALPHA,GL_ONE_MINUS_SRC_ALPHA);
		glEnable(GL_POINT_SMOOTH);
		glEnable(GL_LINE_SMOOTH);
#endif	/* AA_ENABLED */
		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;

}
