#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <cairo.h>
#include <glib.h>
#include <assert.h>
#ifdef _OPENMP
#include <omp.h>
#endif

#include "vector2d.h"
#include "simulator2d.h"

/*
 * TODO: implement text in/out functions.
 */

/* Enable OpneMP only if iteration count exceeds this number. */
#define OMP_THRESHOLD (100)
// Should be larger than 2.0
#define LINK_CELL_FACTOR (2.1)

#define min(a,b) ((a)<(b) ? (a) : (b))
#define max(a,b) ((a)<(b) ? (b) : (a))
#define inertia(p) (0.4*p->mass*p->size*p->size)

#define link_cell_nx(s2d) (ceil(s2d->valid_width / s2d->largest_size / LINK_CELL_FACTOR))
#define link_cell_ny(s2d) (ceil(s2d->valid_height / s2d->largest_size / LINK_CELL_FACTOR))

static void error_exit(const char *msg){
	fprintf(stderr, "E: %s\n", msg);
	exit(1);
}

static void link_cell_free(simulator2d *s2d);

particle2d *s2d_particle_new(){
	particle2d *p = malloc(sizeof(particle2d));
	p->type = NORMAL;
	p->color_r = 1.0;
	p->color_g = 1.0;
	p->color_b = 1.0;
	p->size = 1.0e-3;
	p->mass = 1.0e-3;
	p->r.x = p->r.y = 0.0;
	p->dr.x = p->dr.y = 0.0;
	p->ddr.x = p->ddr.y = 0.0;
	p->dddr.x = p->dddr.y = 0.0;
	p->force.x = p->force.y = 0.0;
	p->theta = p->dtheta = p->ddtheta = p->dddtheta = p->torque = 0.0;
	p->external_force = NULL;
	p->contact_force = NULL;
	p->extra_energy = NULL;
	p->state_updater = NULL;
	p->id_free = NULL;
	return p;
}

void s2d_particle_free(particle2d *p){
	if(p->id_free != NULL)
		p->id_free(p);
	free(p);
}

simulator2d *s2d_new(){
	simulator2d *s2d = malloc(sizeof(simulator2d));
	s2d->p_normal = g_array_new(FALSE, FALSE, sizeof(particle2d*));
	s2d->p_wall = g_array_new(FALSE, FALSE, sizeof(particle2d*));
	s2d->bgcolor_r = 1.0;
	s2d->bgcolor_g = 1.0;
	s2d->bgcolor_b = 1.0;
	s2d->draw_x = 0.0;
	s2d->draw_y = 0.0;
	s2d->zoom = 1000;
	s2d->draw_width = 1200;
	s2d->draw_height = 800;
	s2d->font_size = 15.0;
	s2d->time = 0.0;
	s2d->valid_x = s2d->valid_y = 0.0;
	s2d->valid_width = s2d->valid_height = 1.0;
	s2d->smallest_size = DBL_MAX;
	s2d->largest_size = 0.0;
	s2d->link_cell = NULL;
	s2d->link_cell_nx = -1;
	s2d->link_cell_ny = -1;
	return s2d;
}

void s2d_free(simulator2d *s2d){
	for(int i=s2d->p_normal->len-1; i>=0; i--)
		s2d_remove_particle(s2d, NORMAL, i);
	g_array_free(s2d->p_normal, TRUE);
	for(int i=s2d->p_wall->len-1; i>=0; i--)
		s2d_remove_particle(s2d, WALL, i);
	g_array_free(s2d->p_wall, TRUE);
	link_cell_free(s2d);
	free(s2d);
}

void s2d_add_particle(simulator2d *s2d, particle2d *p){
	switch(p->type){
		case NORMAL:
			s2d->p_normal = g_array_append_val(s2d->p_normal, p);
			break;
		case WALL:
			s2d->p_wall = g_array_append_val(s2d->p_wall, p);
			break;
		default:
			error_exit("the type of a particle is invalid.");
	}
	s2d->smallest_size = min(s2d->smallest_size, p->size);
	s2d->largest_size = max(s2d->largest_size, p->size);
}

void s2d_remove_particle(simulator2d *s2d, particle2d_type type, int index){
	if(type == NORMAL){
		s2d_particle_free(g_array_index(s2d->p_normal, particle2d*, index));
		s2d->p_normal = g_array_remove_index(s2d->p_normal, index);
	}else if(type == WALL){
		s2d_particle_free(g_array_index(s2d->p_wall, particle2d*, index));
		s2d->p_wall = g_array_remove_index(s2d->p_wall, index);
	}else{
		error_exit("the type of a particle is invalid.");
	}
}

static void link_cell_free(simulator2d *s2d){
	if(s2d->link_cell_nx < 0 || s2d->link_cell_ny < 0)
		return;
	for(int x=0; x<s2d->link_cell_nx; x++){
		for(int y=0; y<s2d->link_cell_ny; y++)
			g_array_free(s2d->link_cell[x][y], TRUE);
		free(s2d->link_cell[x]);
	}
	free(s2d->link_cell);
	s2d->link_cell_nx = s2d->link_cell_ny = -1;
}

static void check_valid_area(simulator2d *s2d){
	double valid_l = DBL_MAX;
	double valid_r = -DBL_MAX;
	double valid_u = -DBL_MAX;
	double valid_d = DBL_MAX;
	for(int i=0; i<s2d->p_normal->len; i++){
		particle2d *p = g_array_index(s2d->p_normal, particle2d*, i);
		valid_l = min(valid_l, p->r.x);
		valid_r = max(valid_r, p->r.x);
		valid_u = max(valid_u, p->r.y);
		valid_d = min(valid_d, p->r.y);
	}
	for(int i=0; i<s2d->p_wall->len; i++){
		particle2d *p = g_array_index(s2d->p_wall, particle2d*, i);
		valid_l = min(valid_l, p->r.x);
		valid_r = max(valid_r, p->r.x);
		valid_u = max(valid_u, p->r.y);
		valid_d = min(valid_d, p->r.y);
	}
	s2d->valid_x = valid_l;
	s2d->valid_y = valid_d;
	s2d->valid_width = valid_r - valid_l;
	s2d->valid_height = valid_u - valid_d;
}

static void make_link_cell(simulator2d *s2d){
	check_valid_area(s2d);
	int nx = link_cell_nx(s2d);
	int ny = link_cell_ny(s2d);
	if(s2d->link_cell_nx < nx || s2d->link_cell_ny < ny){
		link_cell_free(s2d);
		s2d->link_cell = malloc(sizeof(GArray**) * nx);
		for(int x=0; x<nx; x++){
			s2d->link_cell[x] = malloc(sizeof(GArray*) * ny);
			for(int y=0; y<ny; y++)
				s2d->link_cell[x][y] = g_array_new(FALSE, FALSE, sizeof(particle2d*));
		}
		s2d->link_cell_nx = nx;
		s2d->link_cell_ny = ny;
	}else{
		#ifdef _OPENMP
		#pragma omp parallel for if(nx>OMP_THRESHOLD)
		#endif
		for(int x=0; x<nx; x++)
			for(int y=0; y<ny; y++)
				if(s2d->link_cell[x][y]->len > 0)
					s2d->link_cell[x][y] = g_array_remove_range(s2d->link_cell[x][y], 0, s2d->link_cell[x][y]->len);
	}
	for(int i=0; i<s2d->p_normal->len; i++){
		particle2d *p = g_array_index(s2d->p_normal, particle2d*, i);
		int x = floor((p->r.x - s2d->valid_x) / s2d->largest_size / LINK_CELL_FACTOR);
		int y = floor((p->r.y - s2d->valid_y) / s2d->largest_size / LINK_CELL_FACTOR);
		assert(0 <= x && x < nx);
		assert(0 <= y && y < ny);
		s2d->link_cell[x][y] = g_array_append_val(s2d->link_cell[x][y], p);
	}
	for(int i=0; i<s2d->p_wall->len; i++){
		particle2d *p = g_array_index(s2d->p_wall, particle2d*, i);
		int x = floor((p->r.x - s2d->valid_x) / s2d->largest_size / LINK_CELL_FACTOR);
		int y = floor((p->r.y - s2d->valid_y) / s2d->largest_size / LINK_CELL_FACTOR);
		assert(0 <= x && x < nx);
		assert(0 <= y && y < ny);
		s2d->link_cell[x][y] = g_array_append_val(s2d->link_cell[x][y], p);
	}
}

static void add_force(simulator2d *s2d, particle2d *p, int x, int y){
	if(p->contact_force == NULL)
		return;
	if(x < 0 || y < 0 || x >= link_cell_nx(s2d) || y >= link_cell_ny(s2d))
		return;
	for(int i=0; i<s2d->link_cell[x][y]->len; i++){
		particle2d *q = g_array_index(s2d->link_cell[x][y], particle2d*, i);
		if(p->type == WALL && q->type == WALL)
			continue;
		p->contact_force(p, q);
	}
}

void s2d_evolute(simulator2d *s2d, double dt){
	//
	// predictor
	#ifdef _OPENMP
	#pragma omp parallel if(s2d->p_normal->len > OMP_THRESHOLD || s2d->p_wall->len > OMP_THRESHOLD)
	{
	#pragma omp for
	#endif
	for(int i=0; i<s2d->p_normal->len; i++){
		particle2d *p = g_array_index(s2d->p_normal, particle2d*, i);
		p->r = v2d_add(
			v2d_add(
				p->r,
				v2d_multiply(dt, p->dr)
			),
			v2d_add(
				v2d_multiply(dt*dt/2.0, p->ddr),
				v2d_multiply(dt*dt*dt/6.0, p->dddr)
			)
		);
		p->dr = v2d_add(v2d_add(p->dr, v2d_multiply(dt, p->ddr)), v2d_multiply(dt*dt/2.0, p->dddr));
		p->ddr = v2d_add(p->ddr, v2d_multiply(dt, p->dddr));
		p->theta += dt*p->dtheta + 0.5*dt*dt*p->ddtheta + (1.0/6.0)*dt*dt*dt*p->dddtheta;
		p->dtheta += dt*p->ddtheta + 0.5*dt*dt*p->dddtheta;
		p->ddtheta += dt*p->dddtheta;
	}
	//
	// force evaluation (link cell)
	#ifdef _OPENMP
	#pragma omp single
	#endif
	make_link_cell(s2d);
	int nx = link_cell_nx(s2d);
	int ny = link_cell_ny(s2d);
	// initialize force and torque
	#ifdef _OPENMP
	#pragma omp for
	#endif
	for(int i=0; i<s2d->p_normal->len; i++){
		particle2d *p = g_array_index(s2d->p_normal, particle2d*, i);
		p->force.x = p->force.y = 0.0;
		p->torque = 0.0;
	}
	#ifdef _OPENMP
	#pragma omp for
	#endif
	for(int i=0; i<s2d->p_wall->len; i++){
		particle2d *p = g_array_index(s2d->p_wall, particle2d*, i);
		p->force.x = p->force.y = 0.0;
		p->torque = 0.0;
	}
	// external force
	#ifdef _OPENMP
	#pragma omp for
	#endif
	for(int i=0; i<s2d->p_normal->len; i++){
		particle2d *p = g_array_index(s2d->p_normal, particle2d*, i);
		if(p->external_force != NULL)
			p->external_force(p);
	}
	// contact force
	#ifdef _OPENMP
	#pragma omp for
	#endif
	for(int x=0; x<nx; x++)
		for(int y=0; y<ny; y++)
			for(int i=0; i<s2d->link_cell[x][y]->len; i++){
				particle2d *p = g_array_index(s2d->link_cell[x][y], particle2d*, i);
				if(p->contact_force == NULL)
					continue;
				for(int j=i+1; j<s2d->link_cell[x][y]->len; j++)
					p->contact_force(p, g_array_index(s2d->link_cell[x][y], particle2d*, j));
				add_force(s2d, p, x-1, y+1);
				add_force(s2d, p, x  , y+1);
				add_force(s2d, p, x+1, y+1);
				add_force(s2d, p, x+1, y);
			}
	//
	// corrector
	#ifdef _OPENMP
	#pragma omp for
	#endif
	for(int i=0; i<s2d->p_normal->len; i++){
		particle2d *p = g_array_index(s2d->p_normal, particle2d*, i);
		vector2d delta_r = v2d_subtract(v2d_divide(p->mass, p->force), p->ddr);
		p->r = v2d_add(p->r, v2d_multiply(dt*dt/12.0, delta_r));
		p->dr = v2d_add(p->dr, v2d_multiply(dt*(5.0/12.0), delta_r));
		p->ddr = v2d_add(p->ddr, delta_r);
		p->dddr = v2d_add(p->dddr, v2d_divide(dt, delta_r));
		double delta_theta = p->torque / inertia(p) - p->ddtheta;
		p->theta += dt*dt/12.0 * delta_theta;
		p->dtheta += dt*(5.0/12.0) * delta_theta;
		p->ddtheta += delta_theta;
		p->dddtheta += delta_theta / dt;
	}
	//
	// call state_updater
	#ifdef _OPENMP
	#pragma omp for
	#endif
	for(int i=0; i<s2d->p_wall->len; i++){
		particle2d *p = g_array_index(s2d->p_wall, particle2d*, i);
		if(p->state_updater != NULL)
			p->state_updater(p, s2d->time);
	}
	#ifdef _OPENMP
	#pragma omp for
	#endif
	for(int i=0; i<s2d->p_normal->len; i++){
		particle2d *p = g_array_index(s2d->p_normal, particle2d*, i);
		if(p->state_updater != NULL)
			p->state_updater(p, s2d->time);
	}
	#ifdef _OPENMP
	}
	#endif
	//
	// add dt to time
	s2d->time += dt;
}

static void write_text_particle(particle2d *p, FILE *f){
	fprintf(
		f, "%d %e %e %e %e %e %e %e %e %e %e %e %e %e %e\n",
		p->type, p->size, p->mass,
		p->r.x, p->r.y, p->dr.x, p->dr.y, p->ddr.x, p->ddr.y, p->dddr.x, p->dddr.y,
		p->theta, p->dtheta, p->ddtheta, p->dddtheta
	);
}
void s2d_write_text(const simulator2d *s2d, const char *fname){
	FILE *f;
	f = fopen(fname, "w");
	if(f == NULL)
		error_exit("failed to write text.");
	for(int i=0; i<s2d->p_normal->len; i++){
		particle2d *p = g_array_index(s2d->p_normal, particle2d*, i);
		write_text_particle(p, f);
	}
	for(int i=0; i<s2d->p_wall->len; i++){
		particle2d *p = g_array_index(s2d->p_wall, particle2d*, i);
		write_text_particle(p, f);
	}
	fclose(f);
}

static void draw_particle(cairo_t *cr, const simulator2d *s2d, const particle2d *p){
	const double ix = s2d->zoom * (p->r.x - s2d->draw_x);
	const double iy = s2d->draw_height - s2d->zoom * (p->r.y - s2d->draw_y);
	const double ir = s2d->zoom * p->size;
	// draw position
	cairo_save(cr);
	cairo_set_source_rgb(cr, p->color_r, p->color_g, p->color_b);
	cairo_arc(
		cr, ix, iy, ir,
		0.0, 2.0*M_PI
	);
	cairo_fill(cr);
	cairo_restore(cr);
	if(p->type == WALL)
		return;
	cairo_save(cr);
	// draw direction
	cairo_set_line_width(cr, 1.5);
	cairo_set_source_rgb(cr, s2d->bgcolor_r, s2d->bgcolor_g, s2d->bgcolor_b);
	double s = 0.4;
	cairo_move_to(cr, ix + s*ir*sin(p->theta+M_PI/2), iy + s*ir*cos(p->theta+M_PI/2));
	cairo_line_to(cr, ix - ir*sin(p->theta)         , iy - ir*cos(p->theta));
	cairo_line_to(cr, ix + s*ir*sin(p->theta-M_PI/2), iy + s*ir*cos(p->theta-M_PI/2));
	cairo_close_path(cr);
	cairo_fill(cr);
	cairo_restore(cr);
}

void s2d_write_cairo(const simulator2d *s2d, cairo_t *cr){
	// draw background
	cairo_set_source_rgb(cr, s2d->bgcolor_r, s2d->bgcolor_g, s2d->bgcolor_b);
	cairo_rectangle(cr, 0, 0, s2d->draw_width, s2d->draw_height);
	cairo_fill(cr);
	// draw particles
	cairo_set_line_width(cr, 0.0);
	for(int i=0; i<s2d->p_normal->len; i++){
		particle2d *p = g_array_index(s2d->p_normal, particle2d*, i);
		draw_particle(cr, s2d, p);
	}
	for(int i=0; i<s2d->p_wall->len; i++){
		particle2d *p = g_array_index(s2d->p_wall, particle2d*, i);
		draw_particle(cr, s2d, p);
	}
	// draw informations
	char str[1024];
	cairo_set_source_rgb(cr, 0.0, 0.0, 0.0);
	cairo_select_font_face(cr, "sans-serif", CAIRO_FONT_SLANT_NORMAL, CAIRO_FONT_WEIGHT_NORMAL);
	cairo_set_font_size(cr, s2d->font_size);
	sprintf(str, "time: %f sec", s2d->time);
	cairo_move_to(cr, 0.5*s2d->font_size, 1.1*s2d->font_size);
	cairo_show_text(cr, str);
	// calculate energy
	double energy = 0.0;
	for(int i=0; i<s2d->p_normal->len; i++){
		particle2d *p = g_array_index(s2d->p_normal, particle2d*, i);
		energy += 0.5 * p->mass * v2d_abs2(p->dr) + 0.5 * inertia(p) * p->dtheta * p->dtheta;
		energy += p->extra_energy(p);
	}
	sprintf(str, "energy: %f J", energy);
	cairo_move_to(cr, 0.5*s2d->font_size, 2.2*s2d->font_size);
	cairo_show_text(cr, str);
}

void s2d_write_png(const simulator2d *s2d, const char *fname){
	// initialize cairo
	cairo_surface_t *surface;
	cairo_t *cr;
	surface = cairo_image_surface_create(CAIRO_FORMAT_RGB24, s2d->draw_width, s2d->draw_height);
	cr = cairo_create(surface);
	//write
	s2d_write_cairo(s2d, cr);
	// write out PNG
	if(cairo_surface_write_to_png(surface, fname) != CAIRO_STATUS_SUCCESS)
		error_exit("failed to write out PNG image.");
	// free memory
	cairo_destroy(cr);
	cairo_surface_destroy(surface);
}